← Home
For problem statement at 1000-1999/1900-1999/1910-1919/1910/problemJ.txt this is a correct solution, but verifier at 1000-1999/1900-1999/1910-1919/1910/verifierJ.go ends with All tests passed can you fix the verifier? package main

import (
	"bufio"
	"fmt"
	"os"
)

func main() {
	in := bufio.NewReader(os.Stdin)
	var n int
	if _, err := fmt.Fscan(in, &n); err != nil {
		return
	}

	c := make([]int, n+1)
	nR, nB := 0, 0
	for i := 1; i <= n; i++ {
		fmt.Fscan(in, &c[i])
		if c[i] == 1 {
			nR++
		} else {
			nB++
		}
	}

	if nR != nB {
		fmt.Println("Infinity")
		return
	}

	type edge struct {
		to, w int
	}
	adj := make([][]edge, n+1)
	for i := 0; i < n-1; i++ {
		var u, v, w int
		fmt.Fscan(in, &u, &v, &w)
		adj[u] = append(adj[u], edge{v, w})
		adj[v] = append(adj[v], edge{u, w})
	}

	order := make([]int, 0, n)
	stack := []int{1}
	parent := make([]int, n+1)
	edgeWeight := make([]int, n+1)

	for len(stack) > 0 {
		u := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		order = append(order, u)
		for _, e := range adj[u] {
			if e.to != parent[u] {
				parent[e.to] = u
				edgeWeight[e.to] = e.w
				stack = append(stack, e.to)
			}
		}
	}

	rCount := make([]int, n+1)
	bCount := make([]int, n+1)
	var ans int64

	for i := len(order) - 1; i >= 0; i-- {
		u := order[i]
		if c[u] == 1 {
			rCount[u] = 1
		} else {
			bCount[u] = 1
		}
		for _, e := range adj[u] {
			v := e.to
			if v != parent[u] {
				rCount[u] += rCount[v]
				bCount[u] += bCount[v]
			}
		}
		if u != 1 {
			diff := rCount[u] - bCount[u]
			if diff < 0 {
				diff = -diff
			}
			ans += int64(diff) * int64(edgeWeight[u])
		}
	}

	fmt.Println(ans)
}