树的存储：父亲表示法与孩子表示法

        
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

        # 树的存储：父亲表示法（parent / 双亲表示法）
# 每个点只记录自己的父亲 fa[u]，根的父亲记为 0。
# 信息方向：向上走一步 O(1)；向下找一个孩子要扫描全部点 O(n)。
# 适用：自底向上递推、按父亲跳祖先、只关心"向上"信息的问题。
# 如何使用：
# 1. pt = ParentTree(n, root)
# 2. pt.set_fa(v, u)              # u 是 v 的父亲，即 fa[v] = u
# 3. pt.for_each_ancestor(v, f)   # 从 v 向上依次访问所有真祖先
# 4. pt.depth(v)                  # v 的深度，根的深度为 0

from collections.abc import Callable


class ParentTree:
    """fa[v] 是 v 的父亲；fa[root] = 0，根没有父亲。"""

    n: int
    root: int
    fa: list[int]

    def __init__(self, n: int, root: int = 1) -> None:
        self.n = n
        self.root = root
        self.fa = [0] * (n + 1)  # 下标 1..n

    def set_fa(self, v: int, f: int) -> None:
        """把 v 的父亲设成 f；注意参数顺序和 fa[v] = f 一致。"""
        self.fa[v] = f

    def is_root(self, u: int) -> bool:
        return self.fa[u] == 0

    def for_each_ancestor(self, v: int, func: Callable[[int], None]) -> None:
        """从 v 向上依次访问所有真祖先（不含 v 自己），一路走到根。"""
        u = self.fa[v]
        while u != 0:
            func(u)
            u = self.fa[u]

    def depth(self, v: int) -> int:
        """v 的深度：从 v 走到根要经过多少条边（树中路径唯一，步数确定）。"""
        d = 0
        u = self.fa[v]
        while u != 0:
            d += 1
            u = self.fa[u]
        return d
    

        
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

        # 树的存储：孩子表示法
# 每个点记录自己的孩子列表 ch[u]，顺手记下父亲 fa[u]，向上向下就都 O(1) 了。
# 枚举 u 的全部孩子只要 O(孩子数)；空间 O(n)，所有列表一共只装 n-1 个孩子。
# 适用：树形 DP、DFS/BFS 遍历、子树统计等"向下走"的问题。
# 如何使用：
# 1. ct = ChildTree(n, root)
# 2. ct.add_child(u, v)       # u 是 v 的父亲
# 3. ct.for_each_child(u, f)  # 遍历 u 的所有孩子
# 4. ct.fa[u]                 # 向上找 u 的父亲

from collections.abc import Callable


class ChildTree:
    """ch[u] 是 u 的孩子列表（按加入顺序）；fa[u] 是 u 的父亲，fa[root] = 0。"""

    n: int
    root: int
    ch: list[list[int]]
    fa: list[int]

    def __init__(self, n: int, root: int = 1) -> None:
        self.n = n
        self.root = root
        self.ch = [[] for _ in range(n + 1)]  # 下标 1..n
        self.fa = [0] * (n + 1)

    def add_child(self, u: int, v: int) -> None:
        """u 是 v 的父亲：把 v 挂到 u 的孩子列表末尾。"""
        self.ch[u].append(v)
        self.fa[v] = u

    def for_each_child(self, u: int, func: Callable[[int], None]) -> None:
        """遍历 u 的所有孩子。"""
        for v in self.ch[u]:
            func(v)
    

        
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

        # 完整实现：从无向边输入建一棵有根树，输出每个点的深度和子树大小
# 输入：
#   第一行 n root
#   接下来 n-1 行 u v（无向边）
# 输出：
#   第一行 depth[1..n]，第二行 size[1..n]
#
# 两种表示法配合使用：
#   孩子表示法 ch[]：自顶向下求深度
#   父亲表示法 fa[]：按 DFS 逆序自底向上累加子树大小

import sys

def main() -> None:
    it = iter(sys.stdin.read().split())
    n = int(next(it))
    root = int(next(it))

    # 第一步：树是特殊的图，先用邻接表存 n-1 条无向边
    g: list[list[int]] = [[] for _ in range(n + 1)]
    for _ in range(n - 1):
        u = int(next(it))
        v = int(next(it))
        g[u].append(v)
        g[v].append(u)

    ch: list[list[int]] = [[] for _ in range(n + 1)]  # 孩子表示法
    fa: list[int] = [0] * (n + 1)                     # 父亲表示法
    depth: list[int] = [0] * (n + 1)                  # 深度：到根的路径上的边数
    sz: list[int] = [1] * (n + 1)                     # 子树大小，先算上自己

    # 从 root 出发 DFS（用栈模拟），确定每个点的父亲和孩子，同时得到遍历顺序
    order: list[int] = []  # DFS 顺序：父亲一定排在孩子前面
    stk = [root]
    fa[root] = 0  # 根没有父亲
    while stk:
        u = stk.pop()
        order.append(u)
        for v in g[u]:
            if v == fa[u]:  # 树上 u 只有父亲这一条"回头边"
                continue
            fa[v] = u       # 第一次走到 v 的点就是 v 的父亲
            ch[u].append(v)
            depth[v] = depth[u] + 1
            stk.append(v)

    # 自底向上（父亲表示法）：逆序遍历，把子树大小累加到父亲上
    for u in reversed(order):
        if u != root:
            sz[fa[u]] += sz[u]

    print("depth:", *(depth[i] for i in range(1, n + 1)))
    print("size:", *(sz[i] for i in range(1, n + 1)))


if __name__ == "__main__":
    main()