题解 | 计树

from collections import deque
import sys


while True:
    try:
        input = sys.stdin.read
        data = input().split()
        n = int(data[0])
        edges = [[] for _ in range(n + 1)]  # edges[i]表示和节点i相连的节点
        idx = 1
        for _ in range(n - 1):
            u, v = int(data[idx]), int(data[idx + 1])
            idx += 2
            edges[u].append(v)
            edges[v].append(u)

        parent = [0] * (n + 1)  # parent[i]表示节点i的父节点
        children = [[] for _ in range(n + 1)]   # children[i]表示节点i的子节点
        q = deque([1])  # 创建一个双端队列，并初始化为[1]
        parent[1] = 0
        while q:
            u = q.popleft()     # 广度优先遍历，构建树
            for v in edges[u]:
                if v != parent[u]:
                    parent[v] = u
                    children[u].append(v)
                    q.append(v)

        k = int(data[idx])
        idx += 1
        V = list(map(int, data[idx:idx + k]))
        idx += k

        # 标记是否在集合 V 中
        is_in_V = [0] * (n + 1)
        for node in V:
            is_in_V[node] = 1

        # 计算 s 数组
        s = [0] * (n + 1)   # s[i]表示节点i的子树上在集合V中的节点的个数
        stack = [(1, False)]
        while stack:
            node, visited = stack.pop()
            if not visited:
                stack.append((node, True))
                # 逆序压入子节点以保证处理顺序正确
                for child in reversed(children[node]):
                    stack.append((child, False))
            else:
                s_val = is_in_V[node]
                for child in children[node]:
                    s_val += s[child]
                s[node] = s_val

        # 计算 cnt 数组
        cnt = [0] * (n + 1)
        for i in range(1, n + 1):
            total = s[i] * s[i]
            sum_child = sum(s[child] * s[child] for child in children[i])
            cnt[i] = total - sum_child
            # 从节点i的子树上任取2个节点，这两个节点刚好落在两个不同子树的可能情况，即为cnt[i]

        # 输出结果
        print(*cnt[1:])
    except:
        break