统计同值子路径的数量

题目描述

给定一棵树，树上有 \(n\) 个节点，节点编号从 0 到 \(n-1\)，树以无向边列表形式给出。
每个节点有一个整数值 \(value[i]\)。
定义一条同值路径为树中的一条简单路径（不重复经过同一个节点），且路径上所有节点的值相等。
注意，路径可以只包含一个节点（长度为 1）。
请计算这棵树中同值路径的数量。

示例：
n = 5
边列表：edges = [[0,1],[1,2],[1,3],[3,4]]
节点值：values = [2, 2, 2, 2, 2]
树形结构如下：

0 -- 1 -- 2
     |
     3
     |
     4

所有节点值都是 2。
同值路径包括：

• 单个节点：5 条（每个节点自身）
• 长度 2 的路径：边 (0,1), (1,2), (1,3), (3,4) → 4 条
• 长度 3 的路径：0-1-2, 1-3-4 → 2 条
• 长度 4 的路径：0-1-3-4 → 1 条
• 长度 5 的路径：无（因为树是简单路径，不会重复节点）
  总数为 5 + 4 + 2 + 1 = 12 条。

请设计算法，计算任意给定的树和节点值列表，返回同值路径的总数。

解题思路

这是一个树形动态规划（本质是线性 DP 在树上的扩展）问题。
由于树是无环连通图，我们可以以任意节点为根（比如 0），将其转化为有根树，然后自底向上（后序遍历）进行 DP 计算。

步骤拆解

步骤 1：将问题转化为有根树上的递归子问题

设以节点 \(u\) 为根的子树中，计算两个值：

• \(cnt[u]\)：在 \(u\) 的子树中，从 \(u\) 出发向下的同值路径的最大长度（即向下延伸的最长同值路径的边数）。
  注意：这条路径必须从 \(u\) 出发，并且路径上所有节点值与 \(u\) 相同。
  如果 \(u\) 的子节点 \(v\) 的值与 \(u\) 不同，则从 \(u\) 往 \(v\) 方向无法延伸同值路径，贡献为 0。
• 统计以 \(u\) 为最高点（路径的转折点）的同值路径数量。
  这是因为任意一条同值路径，在树中都有一个“最高点”（深度最小的节点）。
  我们可以按最高点来分类计数，不重不漏。

步骤 2：递推关系

假设当前节点 \(u\) 的值是 \(val\)。
考虑它的每个子节点 \(v\)（在 DFS 时避免走回父节点）：

• 如果 \(value[v] = val\)，则从 \(u\) 走到 \(v\) 可以形成同值路径，且从 \(v\) 往下还能走 \(cnt[v]\) 条边。
  所以从 \(u\) 到 \(v\) 方向的同值路径长度是 \(1 + cnt[v]\)。
• 如果 \(value[v] \neq val\)，则这个方向不能延伸同值路径，贡献 0。

计算 \(cnt[u]\)：
\(cnt[u]\) 等于所有子节点方向中，能延伸的同值路径长度的最大值。
即：

\[cnt[u] = \max_{v \in children(u)} \big( [value[v]=value[u]] \cdot (1 + cnt[v]) \big) \]

其中 \([P]\) 表示当条件 P 成立时为 1，否则为 0。

计算以 \(u\) 为最高点的同值路径数量：
这种情况的路径在 \(u\) 处可能向两个不同的子节点方向延伸（形成一个倒 V 形路径），也可能只向一个方向延伸（单链）。
我们可以这样统计：

1. 对于单链路径（只从 \(u\) 向一个方向延伸，包括只有 \(u\) 自身）：
   每个方向（包括自身）如果值与 \(u\) 相同，则从 \(u\) 到该方向的路径长度是 \(1 + cnt[v]\)。
   这个长度就是该方向上以 \(u\) 为一个端点、向这个方向延伸的同值路径数量。
   注意：单个节点 \(u\) 也算一条路径，对应长度 1，这里可以统一到公式中，即每个方向贡献的长度就是路径的节点数。

2. 对于两个方向组合的路径：
   从子节点方向1向下走 \(a\) 条边，从子节点方向2向下走 \(b\) 条边，再加上 \(u\) 自身，总节点数 = \(a + b + 1\)。
   这种路径的数量是 \(a \times b\) 吗？不，因为路径是由两个方向的延伸长度组合而成，但这两个延伸是独立的，所以路径的数量是 \((a) \times (b)\)，其中 \(a\) 和 \(b\) 是两个不同方向的有效延伸边数（每个方向的边数等于从 \(u\) 到该方向最远的同值节点的边数，即 \(cnt[child]\)）。

具体算法：

• 收集所有子节点方向的有效延伸长度列表 \(lengths\)，其中 \(lengths[i] = 1 + cnt[child_i]\) 当且仅当 \(value[child_i] = value[u]\)，否则为 0。
• 以 \(u\) 为最高点的路径总数 = 所有长度之和（单链） + 所有两两长度乘积之和（双链）。

步骤 3：算法流程

1. 建图（邻接表）。

2. DFS 后序遍历，对每个节点 \(u\)：

   • 初始化 \(cnt[u] = 0\)。
   • 收集所有子节点方向的有效延伸长度到数组 \(len\)。
   • 更新 \(cnt[u] = \max(len)\)（如果 \(len\) 为空，则 \(cnt[u]=0\)）。
   • 计算以 \(u\) 为最高点的路径数：

     • 单链路径数 = \(\sum len\)（每个 \(len\) 就是该方向上从 \(u\) 出发的同值路径数，包括自身吗？注意：\(len\) 是从 \(u\) 到该方向最远点的边数，所以路径节点数为 \(len+1\)，但这里我们只统计不同路径。实际上，从 \(u\) 出发的长度为 k 的路径只有 1 条（最远的那条），但我们这里要数的是所有可能的同值路径，不仅仅是“最远”的。
       因此，正确方法是：每个方向的有效延伸长度 \(L = 1 + cnt[v]\) 表示从 \(u\) 出发，向这个方向可以形成 \(L\) 条不同的同值路径（长度分别为 1, 2, ..., L 的边数，对应节点数 2, 3, ..., L+1）。但等一下，这里容易出错，我们重新明确：

       我们需要统计的是：以 \(u\) 为最高点的路径，路径上所有节点值相等。
       路径可以分两类：
       A. 只向一个子节点方向延伸（或只有 u 自身）：
       对于某个子节点方向，如果 \(value[v]=val\)，则从 u 到 v 是一条（u, v），还可以继续走到 v 的子节点，直到最长 cnt[v] 条边。
       所以这个方向能形成的以 u 为一个端点的路径数 = \(1 + cnt[v]\)，因为可以选择延伸 0,1,...,cnt[v] 条边。
       延伸 0 条边就是路径 (u) 自身，但注意 (u) 会被每个方向都算一次？不对，路径 (u) 只有一个，不能重复计算。
       所以更好的方法是：单链路径数 = 所有方向的 (1 + cnt[v]) 之和，然后减去（方向数 - 1）来修正 u 自身重复计数？这很乱。
       更清晰的做法：
       将 u 自身单独算 1 条路径。
       对于每个子方向，如果值相同，这个方向能额外贡献 cnt[v] 条以 u 为一个端点的路径（对应边数 1,2,...,cnt[v]）。
       所以单链路径总数 = 1 + 所有方向的 cnt[v] 之和（当值相同时）。

     • 双链路径数：选择两个不同的方向 i, j，值都与 u 相同，则可以组成路径，路径节点是 u 连接两个方向向下延伸的部分。
       对于方向 i 可以选择延伸 1,...,cnt[child_i] 条边，方向 j 可以选择延伸 1,...,cnt[child_j] 条边。
       组合数 = cnt[child_i] * cnt[child_j]。
       对所有不同方向对求和即可。

3. 累加所有节点的“以 u 为最高点的路径数”得到最终答案。

步骤 4：举例验证

用之前的例子，值全为 2，树形结构：
0-1-2 和 1-3-4。

以节点 1 为根（方便计算），但题目无根，我们任选 0 为根 DFS。

树结构（根 0）：
0 -> 1
1 -> 2, 3
3 -> 4

DFS 后序：
节点 2: cnt=0, 单链:1, 双链:0, 总=1
节点 4: cnt=0, 单链:1, 双链:0, 总=1
节点 3: 子 4 同值，cnt[3]=1, 单链:1+cnt[4]=1, 双链:0, 总=1（这里注意：节点 3 自身 1 条，3-4 是 1 条，但 3-4 是单链，延伸长度 1 对应路径 3-4。但这里我们统计以 3 为最高点的路径：

• 自身：1 条
• 向子节点方向延伸：长度 1 条（3-4）
  所以总数 2 条。
  但按公式：单链 = 1 + cnt[4] = 1+0? 不对，因为 3 自身是单独算的，而方向延伸是 cnt[4] 条（即 3-4 一条），所以总 = 1 + cnt[4] = 2。
  双链：0。

节点 1：子 2 同值，cnt[2]=0，子 3 同值，cnt[3]=1。
单链路径数 = 1 + cnt[2] + cnt[3] = 1+0+1=2（自身 1 条，1-2 一条，1-3-4? 不对，1-3-4 是以 1 为最高点的双链？不，1-3-4 是单链，但 1-3 是长度 1，1-3-4 是长度 2，但我们只算了 1-3-4 吗？这里矛盾。

这说明我们刚才公式有问题。我们重新整理：

修正公式

定义 \(dp[u]\) 表示从 u 向下延伸的同值路径的最长边数（即最大深度，只算向下的）。
当计算以 u 为最高点的路径时：

1. 只包含 u 的路径：1 条。

2. 向一个子方向延伸的路径：对每个子节点 v，若值相同，则可以从 u 到 v，还可以继续延伸 dp[v] 条边。
   这个方向贡献的路径数 = dp[v] + 1（对应延伸 1 条边到 v，2 条边到 v 的子节点，等等，直到 dp[v]+1 条边）。
   但注意这里“条数”是：从 u 出发，向这个方向可以形成 dp[v] + 1 条不同的路径（长度分别为 1,...,dp[v]+1 条边）。
   但路径长度按节点数算，分别是 2,3,...,dp[v]+2 个节点，但这里我们只关心路径数，不是节点数。
   所以单链路径总数 = 1（自身） + 对所有同值子节点 v 的 (dp[v] + 1)。

3. 向两个方向延伸的路径：选择两个同值子节点 v1, v2，从 v1 方向延伸 1,...,dp[v1]+1 条边，从 v2 方向延伸 1,...,dp[v2]+1 条边，组合数 = (dp[v1]+1) * (dp[v2]+1)。

所以以 u 为最高点的同值路径总数 =

\[1 + \sum_{v \in children, val[v]=val[u]} (dp[v]+1) + \sum_{v1 \ne v2, val相同} (dp[v1]+1)(dp[v2]+1) \]

举例验证

节点 2：dp=0，总=1。
节点 4：dp=0，总=1。
节点 3：子 4 同值，dp[4]=0，dp[3]=1（最长边数=1）。
以 3 为最高点：
自身 1 条
向子方向延伸：dp[4]+1=1 条（3-4）
双链：0
总=2 条（3，3-4）。对。

节点 1：子 2 同值 dp=0，子 3 同值 dp=1。
dp[1] = max(0+1, 1+1) = 2（最长边数 2，路径 1-3-4）。
以 1 为最高点：
自身 1 条
向子方向延伸：(0+1)+(1+1)=3 条（分别是 1-2，1-3，1-3-4）
双链：(0+1)(1+1)=12=2 条（分别是 1-2 和 1-3 组成路径 2-1-3，2-1-3-4 不，不对，2-1-3 节点值相同，2-1-3-4 节点值相同，但这是 2 条路径，因为 v2 方向可以选择延伸 1 条边到 3，或 2 条边到 4）
更准确：v1=子2方向，dp=0，延伸可选 1 条边（到2），v2=子3方向，dp=1，延伸可选 1 条边（到3）或 2 条边（到4）。
所以组合：
(1 条边 v1) 和 (1 条边 v2) → 路径 2-1-3
(1 条边 v1) 和 (2 条边 v2) → 路径 2-1-3-4
共 2 条。
所以总数=1+3+2=6 条（以 1 为最高点：1, 1-2, 1-3, 1-3-4, 2-1-3, 2-1-3-4）。

节点 0：子 1 同值 dp=2，dp[0]=3（路径 0-1-3-4）。
以 0 为最高点：
自身 1 条
向子方向延伸：dp[1]+1=2+1=3 条（0-1, 0-1-3, 0-1-3-4）
双链：0（只有一个子节点）
总=4 条。

最后总和：
节点 0:4, 1:6, 2:1, 3:2, 4:1 → 4+6+1+2+1=14 条？但之前我们数出 12 条，哪里多了？
检查：我们数了路径 2-1-3-4，节点值全 2，是合法的，长度 4 条边 5 节点，之前数 12 条时漏了这个。
之前数的 12 条是：
单个节点 5 条，长度 2 的 4 条，长度 3 的 2 条，长度 4 的 1 条（0-1-3-4）→ 5+4+2+1=12。
少了 2-1-3-4 这条长度 4 的路径，所以总数应为 13 条？但我们算得 14 条，还多一条。
多的一条是 0-1-2 吗？0-1-2 长度 3 条边 4 节点，之前数长度 3 的路径是 0-1-2 和 1-3-4 两条。
我们算的以 1 为最高点有 0-1-2 吗？没有，因为 0 是 1 的父节点，不是子节点。以 1 为最高点的路径只能是向下两个子方向，所以 0-1-2 是以 0 为最高点的路径吗？不，0-1-2 的最高点是 0 还是 1？路径 0-1-2 的最高点是 1（深度最小是节点 1），所以它应被算在以 1 为最高点里，但 0 是父节点，我们 DFS 时只考虑子节点方向，所以 0-1-2 不会被算在 1 为最高点，而是算在 0 为最高点。
检查我们的计算：以 0 为最高点，子节点只有 1 同值，所以只有单链，没有双链。单链：0, 0-1, 0-1-3, 0-1-3-4 这 4 条，没有 0-1-2，因为 0-1-2 需要经过 1 到 2，但 2 是 1 的子节点，对 0 来说，2 不是直接子节点，所以 0-1-2 是以 1 为最高点的路径，但 1 的父节点 0 同值，所以 1 向上也可以延伸。
这揭示问题：路径的最高点不一定是路径的一个端点，可能是中间节点，我们之前只考虑当前节点向子节点延伸，但路径可以向上延伸到父节点。
所以我们需要考虑向上延伸的情况，即路径可以穿过当前节点向两个方向延伸（一个向子，一个向父）。

修正为同时考虑父节点方向

我们改为：对每个节点 u，计算以 u 为路径的中间节点（即路径穿过 u，向两个不同方向延伸）的同值路径数。
这样能覆盖所有路径，因为任何长度 ≥ 2 的同值路径都有一个最高节点，这个最高节点就是路径中深度最小的节点。
在 DFS 中，对节点 u，我们只考虑 u 的子节点方向组合（不包括父节点方向），因为父节点方向在 u 的父节点处会考虑到。
这样能保证不重复计数。

所以算法：
DFS 返回 dp[u] 表示从 u 向下延伸的同值路径的最长边数。
在 DFS 过程中，对节点 u，收集所有同值子节点的 (dp[v]+1) 到列表 L。
以 u 为最高点的路径数 =
1（只有 u 自身） +
sum(L) +
sum_{i<j} L[i]*L[j]。
其中 sum(L) 是向一个子方向延伸的路径数（包括自身到子节点的边），sum_{i<j} 是向两个子方向延伸的路径数。

这样计算，对之前例子：

节点 2: L=[], 总=1。
节点 4: L=[], 总=1。
节点 3: 子 4 同值 dp=0, L=[1], 总=1+1+0=2（自身 3，3-4）。
节点 1: 子 2 同值 dp=0 → L1=1, 子 3 同值 dp=1 → L2=2, L=[1,2]。
sumL=3, 两两乘积=1*2=2。
总=1+3+2=6（自身 1，1-2，1-3，1-3-4，2-1-3，2-1-3-4）。
节点 0: 子 1 同值 dp=2 → L=[3]，总=1+3+0=4（自身 0，0-1，0-1-3，0-1-3-4）。
总和=1+1+2+6+4=14 条。

现在我们检查所有同值路径：
单节点：5 条。
2 节点：4 条 (0,1),(1,2),(1,3),(3,4)。
3 节点：3 条 (0-1-2),(1-3-4),(0-1-3)？注意 0-1-3 是 3 节点吗？节点 0,1,3 是 3 个节点，是。
4 节点：1 条 (0-1-3-4)。
5 节点：1 条 (2-1-3-4)。
总数 5+4+3+1+1=14 条，对上了。

步骤 5：最终算法

1. 建图。
2. DFS(u, parent)：
   • 初始化 dp[u]=0, 初始化列表 L=[]。
   • 对每个邻居 v != parent：
     • 递归计算 DFS(v, u)。
     • 如果 value[v]==value[u]，则把 dp[v]+1 加入 L，并更新 dp[u]=max(dp[u], dp[v]+1)。
   • 计算以 u 为最高点的路径数：
     • 单链路径 = 1 + sum(L)（1 是自身，sum(L) 是向每个子方向延伸的路径数）。
     • 双链路径 = 所有 L 中两两乘积之和（可以用 (sum(L)^2 - sum(L[i]^2))/2 计算）。
     • 总 = 单链路径 + 双链路径。
   • 返回 dp[u]。
3. 累加所有节点的“总”得到答案。

时间复杂度

每个节点访问一次，每个边处理一次，O(n)。
两两乘积计算可以优化：双链路径 = (sum^2 - sum_sq)/2，其中 sum 是 L 的和，sum_sq 是 L 的平方和。
总体 O(n)。

示例代码（Python）

def count_univalue_paths(n, edges, values):
    from collections import defaultdict
    
    graph = defaultdict(list)
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)
    
    total_paths = 0
    
    def dfs(u, parent):
        nonlocal total_paths
        L = []
        for v in graph[u]:
            if v == parent:
                continue
            child_len = dfs(v, u)
            if values[v] == values[u]:
                L.append(child_len + 1)
        # dp[u] 是 L 的最大值，如果 L 空则是 0
        dp_u = max(L) if L else 0
        
        # 计算以 u 为最高点的路径数
        sumL = sum(L)
        sum_sq = sum(x*x for x in L)
        single_chain = 1 + sumL
        double_chain = (sumL * sumL - sum_sq) // 2
        total_paths += single_chain + double_chain
        
        return dp_u
    
    dfs(0, -1)
    return total_paths

# 测试
n = 5
edges = [[0,1],[1,2],[1,3],[3,4]]
values = [2,2,2,2,2]
print(count_univalue_paths(n, edges, values))  # 输出 14

这样我们就完成了统计树中同值路径数量的动态规划解法。