For problem statement at 2000-2999/2100-2199/2150-2159/2155/problemF.txt this is a correct solution, but verifier at 2000-2999/2100-2199/2150-2159/2155/verifierF.go ends with All 202 tests passed can you fix the verifier? ```go
package main
import (
"bufio"
"fmt"
"os"
)
// Fast IO
type fastScanner struct {
buf []byte
pos int
r *bufio.Reader
}
func newFastScanner(r *bufio.Reader) *fastScanner {
return &fastScanner{buf: make([]byte, 0, 1<<16), r: r}
}
func (s *fastScanner) read() {
s.buf = s.buf[:cap(s.buf)]
n, _ := s.r.Read(s.buf)
s.buf = s.buf[:n]
s.pos = 0
}
func (s *fastScanner) nextInt() int {
for {
if s.pos >= len(s.buf) {
s.read()
if len(s.buf) == 0 {
return 0
}
}
c := s.buf[s.pos]
if c > ' ' {
break
}
s.pos++
}
res := 0
for {
if s.pos >= len(s.buf) {
s.read()
if len(s.buf) == 0 {
break
}
}
c := s.buf[s.pos]
if c <= ' ' {
break
}
res = res*10 + int(c-'0')
s.pos++
}
return res
}
func main() {
reader := bufio.NewReader(os.Stdin)
writer := bufio.NewWriter(os.Stdout)
defer writer.Flush()
scanner := newFastScanner(reader)
t := scanner.nextInt()
for i := 0; i < t; i++ {
solve(scanner, writer)
}
}
func solve(scanner *fastScanner, writer *bufio.Writer) {
n := scanner.nextInt()
k := scanner.nextInt()
s := scanner.nextInt()
q := scanner.nextInt()
// Adjacency list for the tree
adj := make([][]int, n+1)
for i := 0; i < n-1; i++ {
u := scanner.nextInt()
v := scanner.nextInt()
adj[u] = append(adj[u], v)
adj[v] = append(adj[v], u)
}
// Colors
nodesOfColor := make([][]int, k+1)
nodeColorCount := make([]int, n+1)
for i := 0; i < s; i++ {
v := scanner.nextInt()
x := scanner.nextInt()
nodesOfColor[x] = append(nodesOfColor[x], v)
nodeColorCount[v]++
}
// Queries
type Query struct {
u, v, id int
}
rawQueries := make([]Query, q)
for i := 0; i < q; i++ {
u := scanner.nextInt()
v := scanner.nextInt()
if u > v {
u, v = v, u
}
rawQueries[i] = Query{u, v, i}
}
// BFS for parents
parent := make([]int, n+1)
queue := make([]int, 0, n)
queue = append(queue, 1)
visitedBFS := make([]bool, n+1)
visitedBFS[1] = true
head := 0
for head < len(queue) {
u := queue[head]
head++
for _, v := range adj[u] {
if !visitedBFS[v] {
visitedBFS[v] = true
parent[v] = u
queue = append(queue, v)
}
}
}
// Identify Heavy Nodes
B := 0
for B*B < s {
B++
}
if B < 1 {
B = 1
}
heavyID := make([]int, n+1)
for i := range heavyID {
heavyID[i] = -1
}
var heavyNodes []int
for u := 1; u <= n; u++ {
if nodeColorCount[u] > B {
heavyID[u] = len(heavyNodes)
heavyNodes = append(heavyNodes, u)
}
}
numHeavy := len(heavyNodes)
// Matrix for heavy-heavy queries
heavyQueries := make([][][]int, numHeavy)
for i := 0; i < numHeavy; i++ {
heavyQueries[i] = make([][]int, numHeavy)
}
type LightQuery struct {
v, id int
}
lightQueries := make([][]LightQuery, n+1)
// Distribute queries
for _, qry := range rawQueries {
u, v, id := qry.u, qry.v, qry.id
uHeavy := heavyID[u] != -1
vHeavy := heavyID[v] != -1
if uHeavy && vHeavy {
hu, hv := heavyID[u], heavyID[v]
if hu > hv {
hu, hv = hv, hu
}
heavyQueries[hu][hv] = append(heavyQueries[hu][hv], id)
} else {
// Assign to light node (or the one with fewer colors)
target := u
other := v
if uHeavy {
target = v
other = u
} else if !vHeavy {
if nodeColorCount[u] > nodeColorCount[v] {
target = v
other = u
}
}
lightQueries[target] = append(lightQueries[target], LightQuery{other, id})
}
}
ans := make([]int, q)
// Helper arrays for DSU and Components
dsuParent := make([]int, n+1)
active := make([]bool, n+1)
inComp := make([]bool, n+1)
compNodes := make([][]int, n+1)
leaders := make([]int, 0, n)
var find func(int) int
find = func(i int) int {
if dsuParent[i] == i {
return i
}
dsuParent[i] = find(dsuParent[i])
return dsuParent[i]
}
union := func(i, j int) {
rootI := find(i)
rootJ := find(j)
if rootI != rootJ {
dsuParent[rootI] = rootJ
}
}
for c := 1; c <= k; c++ {
nodes := nodesOfColor[c]
if len(nodes) == 0 {
continue
}
// Prepare nodes
for _, u := range nodes {
dsuParent[u] = u
active[u] = true
}
// Build DSU
for _, u := range nodes {
p := parent[u]
if p != 0 && active[p] {
union(u, p)
}
}
// Group components
currentLeaders := leaders[:0]
for _, u := range nodes {
root := find(u)
if len(compNodes[root]) == 0 {
currentLeaders = append(currentLeaders, root)
}
compNodes[root] = append(compNodes[root], u)
}
// Process components
for _, root := range currentLeaders {
comp := compNodes[root]
for _, u := range comp {
inComp[u] = true
}
// Light queries
for _, u := range comp {
for _, lq := range lightQueries[u] {
if inComp[lq.v] {
ans[lq.id]++
}
}
}
// Heavy queries
var hInComp []int
for _, u := range comp {
if heavyID[u] != -1 {
hInComp = append(hInComp, u)
}
}
for i := 0; i < len(hInComp); i++ {
u := hInComp[i]
hu := heavyID[u]
for j := i; j < len(hInComp); j++ {
v := hInComp[j]
hv := heavyID[v]
h1, h2 := hu, hv
if h1 > h2 {
h1, h2 = h2, h1
}
if len(heavyQueries[h1][h2]) > 0 {
for _, qID := range heavyQueries[h1][h2] {
ans[qID]++
}
}
}
}
for _, u := range comp {
inComp[u] = false
}
compNodes[root] = compNodes[root][:0]
}
for _, u := range nodes {
active[u] = false
}
leaders = currentLeaders
}
for i := 0; i < q; i++ {
if i > 0 {
writer.WriteByte(' ')
}
fmt.Fprintf(writer, "%d", ans[i])
}
writer.WriteByte('\n')
}
```