For problem statement at 1000-1999/1700-1799/1740-1749/1746/problemD.txt this is a correct solution, but verifier at 1000-1999/1700-1799/1740-1749/1746/verifierD.go ends with All tests passed can you fix the verifier? package main
import (
"io"
"os"
"sort"
"strconv"
)
type Frame struct {
u int
x int64
stage uint8
}
func main() {
data, _ := io.ReadAll(os.Stdin)
pos := 0
nextInt := func() int {
for pos < len(data) && (data[pos] < '0' || data[pos] > '9') && data[pos] != '-' {
pos++
}
sign := 1
if data[pos] == '-' {
sign = -1
pos++
}
val := 0
for pos < len(data) && data[pos] >= '0' && data[pos] <= '9' {
val = val*10 + int(data[pos]-'0')
pos++
}
return sign * val
}
t := nextInt()
out := make([]byte, 0, t*24)
for ; t > 0; t-- {
n := nextInt()
k := int64(nextInt())
parent := make([]int, n+1)
deg := make([]int, n+1)
for i := 2; i <= n; i++ {
p := nextInt()
parent[i] = p
deg[p]++
}
children := make([][]int, n+1)
for i := 1; i <= n; i++ {
if deg[i] > 0 {
children[i] = make([]int, 0, deg[i])
}
}
for i := 2; i <= n; i++ {
p := parent[i]
children[p] = append(children[p], i)
}
score := make([]int64, n+1)
for i := 1; i <= n; i++ {
score[i] = int64(nextInt())
}
memoCnt := make([]uint8, n+1)
memoX1 := make([]int64, n+1)
memoV1 := make([]int64, n+1)
memoX2 := make([]int64, n+1)
memoV2 := make([]int64, n+1)
var extra map[uint64]int64
key := func(u int, x int64) uint64 {
return (uint64(uint32(u)) << 32) | uint64(uint32(x))
}
getMemo := func(u int, x int64) (int64, bool) {
if memoCnt[u] > 0 && memoX1[u] == x {
return memoV1[u], true
}
if memoCnt[u] > 1 && memoX2[u] == x {
return memoV2[u], true
}
if extra != nil {
v, ok := extra[key(u, x)]
return v, ok
}
return 0, false
}
setMemo := func(u int, x, v int64) {
if memoCnt[u] == 0 {
memoCnt[u] = 1
memoX1[u] = x
memoV1[u] = v
return
}
if memoX1[u] == x {
memoV1[u] = v
return
}
if memoCnt[u] == 1 {
memoCnt[u] = 2
memoX2[u] = x
memoV2[u] = v
return
}
if memoX2[u] == x {
memoV2[u] = v
return
}
if extra == nil {
extra = make(map[uint64]int64)
}
extra[key(u, x)] = v
}
mustMemo := func(u int, x int64) int64 {
if memoCnt[u] > 0 && memoX1[u] == x {
return memoV1[u]
}
if memoCnt[u] > 1 && memoX2[u] == x {
return memoV2[u]
}
if extra != nil {
return extra[key(u, x)]
}
return 0
}
stack := make([]Frame, 1, n+5)
stack[0] = Frame{u: 1, x: k}
for len(stack) > 0 {
idx := len(stack) - 1
u := stack[idx].u
x := stack[idx].x
if _, ok := getMemo(u, x); ok {
stack = stack[:idx]
continue
}
if stack[idx].stage == 0 {
if x == 0 {
setMemo(u, 0, 0)
stack = stack[:idx]
continue
}
if len(children[u]) == 0 {
setMemo(u, x, x*score[u])
stack = stack[:idx]
continue
}
stack[idx].stage = 1
m := int64(len(children[u]))
q := x / m
r := x % m
if r == 0 {
for _, v := range children[u] {
if _, ok := getMemo(v, q); !ok {
stack = append(stack, Frame{u: v, x: q})
}
}
} else {
qp := q + 1
for _, v := range children[u] {
if _, ok := getMemo(v, q); !ok {
stack = append(stack, Frame{u: v, x: q})
}
if _, ok := getMemo(v, qp); !ok {
stack = append(stack, Frame{u: v, x: qp})
}
}
}
continue
}
m := int64(len(children[u]))
q := x / m
r := int(x % m)
res := x * score[u]
if r == 0 {
for _, v := range children[u] {
res += mustMemo(v, q)
}
} else {
qp := q + 1
deltas := make([]int64, 0, len(children[u]))
for _, v := range children[u] {
a := mustMemo(v, q)
b := mustMemo(v, qp)
res += a
deltas = append(deltas, b-a)
}
sort.Slice(deltas, func(i, j int) bool {
return deltas[i] > deltas[j]
})
for i := 0; i < r; i++ {
res += deltas[i]
}
}
setMemo(u, x, res)
stack = stack[:idx]
}
ans, _ := getMemo(1, k)
out = strconv.AppendInt(out, ans, 10)
out = append(out, '\n')
}
_, _ = os.Stdout.Write(out)
}