← Home
For problem statement at 0-999/700-799/770-779/771/problemC.txt this is a correct solution, but verifier at 0-999/700-799/770-779/771/verifierC.go ends with All tests passed! can you fix the verifier? package main

import (
	"fmt"
	"io"
	"os"
)

func main() {
	input, _ := io.ReadAll(os.Stdin)
	var pos int
	nextInt := func() int {
		for pos < len(input) && input[pos] <= ' ' {
			pos++
		}
		if pos >= len(input) {
			return 0
		}
		res := 0
		for pos < len(input) && input[pos] > ' ' {
			res = res*10 + int(input[pos]-'0')
			pos++
		}
		return res
	}

	n := nextInt()
	if n == 0 {
		return
	}
	k := nextInt()

	adj := make([][]int, n+1)
	for i := 0; i < n-1; i++ {
		u := nextInt()
		v := nextInt()
		adj[u] = append(adj[u], v)
		adj[v] = append(adj[v], u)
	}

	sz := make([]int64, n+1)
	cnt := make([][5]int64, n+1)
	var S [5]int64
	var totDist int64

	var dfs func(u, p int)
	dfs = func(u, p int) {
		sz[u] = 1
		cnt[u][0] = 1
		for _, v := range adj[u] {
			if v == p {
				continue
			}
			dfs(v, u)

			for i := 0; i < k; i++ {
				for j := 0; j < k; j++ {
					S[(i+j+1)%k] += cnt[u][i] * cnt[v][j]
				}
			}

			for j := 0; j < k; j++ {
				cnt[u][(j+1)%k] += cnt[v][j]
			}

			sz[u] += sz[v]
			totDist += sz[v] * (int64(n) - sz[v])
		}
	}

	dfs(1, 0)

	ans := totDist
	for r := 1; r < k; r++ {
		ans += int64(k-r) * S[r]
	}
	ans /= int64(k)

	fmt.Println(ans)
}