For problem statement at 0-999/500-599/510-519/516/problemD.txt this is a correct solution, but verifier at 0-999/500-599/510-519/516/verifierD.go ends with All 204 tests passed can you fix the verifier? package main
import (
"bufio"
"fmt"
"io"
"os"
"sort"
)
type Edge struct {
to int
w int64
}
type Fenwick struct {
n int
tree []int
}
func NewFenwick(n int) *Fenwick {
return &Fenwick{n: n, tree: make([]int, n+2)}
}
func (f *Fenwick) Add(i, v int) {
for i <= f.n {
f.tree[i] += v
i += i & -i
}
}
func (f *Fenwick) Sum(i int) int {
s := 0
for i > 0 {
s += f.tree[i]
i -= i & -i
}
return s
}
func (f *Fenwick) Reset() {
for i := range f.tree {
f.tree[i] = 0
}
}
func build(start int, adj [][]Edge) (int, []int64, []int) {
n := len(adj) - 1
dist := make([]int64, n+1)
for i := 1; i <= n; i++ {
dist[i] = -1
}
parent := make([]int, n+1)
stack := make([]int, 0, n)
stack = append(stack, start)
dist[start] = 0
parent[start] = 0
for len(stack) > 0 {
u := stack[len(stack)-1]
stack = stack[:len(stack)-1]
for _, e := range adj[u] {
if e.to == parent[u] {
continue
}
parent[e.to] = u
dist[e.to] = dist[u] + e.w
stack = append(stack, e.to)
}
}
far := start
for i := 1; i <= n; i++ {
if dist[i] > dist[far] {
far = i
}
}
return far, dist, parent
}
func euler(root int, adj [][]Edge) ([]int, []int) {
n := len(adj) - 1
tin := make([]int, n+1)
tout := make([]int, n+1)
parent := make([]int, n+1)
type Frame struct {
u int
idx int
}
stack := make([]Frame, 0, n)
stack = append(stack, Frame{u: root})
timer := 0
for len(stack) > 0 {
top := &stack[len(stack)-1]
u := top.u
if top.idx == 0 {
timer++
tin[u] = timer
}
if top.idx < len(adj[u]) {
e := adj[u][top.idx]
top.idx++
if e.to == parent[u] {
continue
}
parent[e.to] = u
stack = append(stack, Frame{u: e.to})
} else {
tout[u] = timer
stack = stack[:len(stack)-1]
}
}
return tin, tout
}
func main() {
data, _ := io.ReadAll(os.Stdin)
idx := 0
nextInt := func() int64 {
for idx < len(data) && (data[idx] < '0' || data[idx] > '9') && data[idx] != '-' {
idx++
}
sign := int64(1)
if data[idx] == '-' {
sign = -1
idx++
}
var v int64
for idx < len(data) && data[idx] >= '0' && data[idx] <= '9' {
v = v*10 + int64(data[idx]-'0')
idx++
}
return sign * v
}
n := int(nextInt())
adj := make([][]Edge, n+1)
for i := 0; i < n-1; i++ {
x := int(nextInt())
y := int(nextInt())
v := nextInt()
adj[x] = append(adj[x], Edge{to: y, w: v})
adj[y] = append(adj[y], Edge{to: x, w: v})
}
a, _, _ := build(1, adj)
b, distA, _ := build(a, adj)
_, distB, _ := build(b, adj)
D := distA[b]
dep2 := make([]int64, n+1)
root := 1
for i := 1; i <= n; i++ {
mx := distA[i]
if distB[i] > mx {
mx = distB[i]
}
dep2[i] = 2*mx - D
if dep2[i] < dep2[root] {
root = i
}
}
tin, tout := euler(root, adj)
ids := make([]int, n)
for i := 1; i <= n; i++ {
ids[i-1] = i
}
sort.Slice(ids, func(i, j int) bool {
if dep2[ids[i]] == dep2[ids[j]] {
return ids[i] < ids[j]
}
return dep2[ids[i]] < dep2[ids[j]]
})
depVals := make([]int64, n)
for i, v := range ids {
depVals[i] = dep2[v]
}
q := int(nextInt())
queries := make([]int64, q)
for i := 0; i < q; i++ {
queries[i] = nextInt()
}
bit := NewFenwick(n)
out := bufio.NewWriterSize(os.Stdout, 1<<20)
for _, l := range queries {
bit.Reset()
L2 := 2 * l
best := 0
p := 0
for _, v := range ids {
limit := dep2[v] + L2
for p < n && depVals[p] <= limit {
bit.Add(tin[ids[p]], 1)
p++
}
cnt := bit.Sum(tout[v]) - bit.Sum(tin[v]-1)
if cnt > best {
best = cnt
}
}
fmt.Fprintln(out, best)
}
out.Flush()
}