题解 | 计树

package main

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

func main() {
	in := bufio.NewReader(os.Stdin)
	out := bufio.NewWriter(os.Stdout)
	defer out.Flush()

	var n int
	fmt.Fscan(in, &n)

	// 建图（邻接表）
	graph := make([][]int, n+1)
	for i := 0; i < n-1; i++ {
		var u, v int
		fmt.Fscan(in, &u, &v)
		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	// BFS 建立有根树（以1为根），记录 children
	parent := make([]int, n+1)
	children := make([][]int, n+1)
	visited := make([]bool, n+1)
	queue := []int{1}
	visited[1] = true
	for len(queue) > 0 {
		u := queue[0]
		queue = queue[1:]
		for _, v := range graph[u] {
			if !visited[v] {
				visited[v] = true
				parent[v] = u
				children[u] = append(children[u], v)
				queue = append(queue, v)
			}
		}
	}

	// 读入 k 和集合 V
	var k int
	fmt.Fscan(in, &k)
	isInV := make([]bool, n+1)
	for i := 0; i < k; i++ {
		var x int
		fmt.Fscan(in, &x)
		isInV[x] = true
	}

	// DFS 后序遍历，计算子树中 V 的节点数
	subCnt := make([]int, n+1) // subCnt[u] = 子树 u 中属于 V 的节点数
	var dfs func(int)
	dfs = func(u int) {
		if isInV[u] {
			subCnt[u] = 1
		}
		for _, v := range children[u] {
			dfs(v)
			subCnt[u] += subCnt[v]
		}
	}
	dfs(1)

	// 计算答案：ans[u] = subCnt[u]^2 - sum(child subCnt[child]^2)
	ans := make([]int, n+1)
	for u := 1; u <= n; u++ {
		total := subCnt[u] * subCnt[u]
		for _, v := range children[u] {
			total -= subCnt[v] * subCnt[v]
		}
		ans[u] = total
	}

	// 输出
	for i := 1; i <= n; i++ {
		if i > 1 {
			fmt.Fprint(out, " ")
		}
		fmt.Fprint(out, ans[i])
	}
	fmt.Fprintln(out)
}