2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个

网易专栏18小时前发布 nxnqh
1 0 0

🤖 AI总结

主题

使用树形DP统计有根树中满足不相邻和整除条件的子集数量。

摘要

文章介绍通过树形DP和卷积合并子节点,统计有根树中不相邻且和整除k的非空子集数量,给出Go、Python、C++实现。

关键信息

  • 1 通过DP状态f0和f1处理选与不选当前节点的情况
  • 2 利用卷积合并子节点方案数
  • 3 时间复杂度O(n*k^2)

2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个节点的父节点由一个数组 parent 给出,根节点的父节点为 -1,其他节点的父节点编号一定小于该节点本身。同时,每个节点上还有一个整数值,存放在数组 nums 中,另外给定一个整数 k。

我们需要统计所有满足以下两个条件的非空节点集合的数量:

1. 集合中所有节点的数值之和能够被 k 整除;

  • 2. 集合中不能同时包含任意一个节点和它的直接父节点,也就是说选出的节点在树中不能有相邻的父子关系。

    最终结果需要对 1000000007 取模后输出。

    n == parent.length == nums.length

    1 <= n <= 1000

    parent[0] == -1

    对于所有的 1 <= i < n:

    0 <= parent[i] < i

    1 <= nums[i] <= 1000000000

    1 <= k <= 100

    parent 表示一棵有效的有根树。

    输入: parent = [-1,0,0,0], nums = [2,1,2,1], k = 3。

    输出: 2。

    解释:

    2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个
    在这里插入图片描述

    有效的子集有:

    {1, 2}:节点 1 和 2 都是节点 0 的子节点,且彼此不直接相连。它们的值之和为 1 + 2 = 3 ,可以被 3 整除。

    {2, 3}:节点 2 和 3 也不相邻。它们的值之和为 2 + 1 = 3 ,可以被 3 整除。

    没有其他子集同时满足两个条件。因此,答案是 2 。

    题目来自力扣3939。

    合并子节点的详细过程

    假设我们已经递归计算了某个子节点 y 的状态fy0fy1,现在要将 y 合并到当前节点 x 的f0f1中。

    1. 更新f0(不选 x)

    此时,由于 x 未被选中,子节点 y可以被选,也可以不被选。因此,从 y 子树中选取的合法集合有两种情况:

    • y 不被选,对应fy0

  • • y 被选,对应fy1

    所以,子节点 y 对整体余数的贡献总和为v[i] = fy0[i] + fy1[i](每种余数 i 的方案数相加)。

    现在,当前已有的不选 x 的方案数为f0(这是已经处理完之前若干个兄弟子树的累计结果)。当我们把 y 的贡献合并进来时,相当于将两个“余数分布”进行卷积:新余数 = (i + j) % k,其中 i 来自子节点 y 贡献的余数,j 来自之前已处理的子树贡献的余数。新的方案数累加到nf0[(i+j)%k]中。

    最后,用nf0替换原有的f0

    2. 更新f1(选 x)

    此时,x 已被选中,那么其直接子节点 y 绝对不能选(因为 y 是 x 的子节点,二者相邻)。因此,y 子树只能提供 y 不被选时的方案,即fy0

    类似地,将fy0与当前已有的f1(已经处理完的兄弟子树)进行卷积,得到新的nf1,并替换原有的f1

    递归过程说明

    • 整棵树通过parent数组构建邻接表,根节点为 0。

  • • 从根节点开始执行深度优先搜索(DFS),递归地处理每个节点。

  • • 每个节点在处理完所有子节点后,返回自己的f0f1给父节点。

  • • 父节点在得到子节点的返回结果后,按照上述规则合并。

    最终答案的计算

    当根节点 0 的递归返回后,我们得到了整棵树的f0f1(分别对应不选根和选根两种全局状态)。

    • 合法的非空集合总数 = (不选根时,余数为 0 的方案数) + (选根时,余数为 0 的方案数)。

  • • 但这两个方案数中都包含了空集(因为初始的f0[0]=1就代表空集,而选根时不可能包含空集,所以只有f0里有空集),所以最后需要减去空集这一种方案

    即答案 =(f0[0] + f1[0] - 1) mod MOD,最后取模得到正整数结果。

    时间复杂度

    • 每个节点在合并其每个子节点时,都需要两层循环分别遍历余数 0 到 k-1,因此每次合并的时间开销为O(k^2)

  • • 树中总共有 n 个节点,每条边对应一次合并操作,边的数量为n-1

  • • 因此总时间复杂度为O(n · k^2)

  • • 在本题限制下,n ≤ 1000k ≤ 100,故最多约1000 × 10000 = 1e7次基本运算,完全可行。

    额外空间复杂度

    • 递归深度最坏情况下为O(n)(例如链状树)。

  • • 在递归栈的每一层,每个节点会保存若干个长度为 k 的数组(f0f1,以及合并时的临时数组),因此递归路径上同时存在的数组总大小约为O(k)乘以递归深度,即O(n · k)

  • • 同时,合并过程中产生的临时数组会在函数返回后自动释放,不会累积。

  • • 因此,额外空间复杂度为O(n · k),在给定范围内(n=1000, k=100)约为1e5级别,内存充足。

    示例验证(以题目输入为例)

    parent = [-1,0,0,0],根为 0,子节点为 1、2、3。

  • nums = [2,1,2,1]k=3

  • • 叶子节点 1、2、3 分别递归返回。

  • • 根 0 合并子节点后,最终统计余数为 0 的方案数(减去空集)得到答案 2,即{1,2}{2,3}两种有效子集,与题意相符。

    总结

    该算法利用树形 DP 巧妙地处理了“不相邻”和“和整除 k”两个约束,通过分情况(选/不选当前节点)以及卷积合并子节点的方式,在O(n·k²)时间内完成了统计。代码实现清晰,适合本题的数据规模。

    Go完整代码如下:

    package main

    import (
    "fmt"
    )

    func countValidSubsets(parent []int, nums []int, k int) int {
    const mod = 1_000_000_007
    n := len(parent)
    g := make([][]int, n)
    for i := 1; i < n; i++ {
    p := parent[i]
    g[p] = append(g[p], i)
    }

    var dfs func(int) ([]int, []int)
    dfs = func(x int) ([]int, []int) {
    f0 := make([]int, k) // f0[i] 表示不选 x 时,子树 x 的子集点权和模 k 为 i 的方案数
    f1 := make([]int, k) // f1[i] 表示选 x 时,子树 x 的子集点权和模 k 为 i 的方案数
    f0[0] = 1
    f1[nums[x]%k] = 1

    for _, y := range g[x] {
    fy0, fy1 := dfs(y)

    // 不选 x,那么 y 可选可不选
    nf0 := make([]int, k)
    for i := range k { // 枚举从子树 y 中选出的点权和模 k 为 i
    v := fy0[i] + fy1[i]
    if v == 0 { // 优化
    continue
    }
    for j, w := range f0 { // 枚举从之前的子树中选出的点权和模 k 为 j
    s := (i + j) % k
    nf0[s] = (nf0[s] + v*w) % mod
    }
    }

    // 选 x,那么 y 不能选
    nf1 := make([]int, k)
    for i, v := range fy0 { // 枚举从子树 y 中选出的点权和模 k 为 i
    if v == 0 { // 优化
    continue
    }
    for j, w := range f1 { // 枚举从 x 以及之前的子树中选出的点权和模 k 为 j
    s := (i + j) % k
    nf1[s] = (nf1[s] + v*w) % mod
    }
    }
    f0, f1 = nf0, nf1
    }

    return f0, f1
    }

    f0, f1 := dfs(0)
    // 恰好被 k 整除即模 k 为 0,注意减去空集的方案数 1
    return (f0[0] + f1[0] - 1 + mod) % mod
    }

    func main() {
    parent := []int{-1, 0, 0, 0}
    nums := []int{2, 1, 2, 1}
    k := 3
    result := countValidSubsets(parent, nums, k)
    fmt.Println(result)
    }

    2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个

    Python完整代码如下:

    # -*-coding:utf-8-*-

    import sys

    def countValidSubsets(parent, nums, k):
    MOD = 10**9 + 7
    n = len(parent)
    # 构建邻接表
    g = [[] for _ in range(n)]
    for i in range(1, n):
    p = parent[i]
    g[p].append(i)

    sys.setrecursionlimit(max(1000000, n * 2 + 10))

    def dfs(x):
    # f0: 不选当前节点 x 时的方案数(按模 k 分类)
    # f1: 选当前节点 x 时的方案数(按模 k 分类)
    f0 = [0] * k
    f1 = [0] * k
    f0[0] = 1 # 空集
    f1[nums[x] % k] = 1 # 只含 x 自身的子集

    for y in g[x]:
    fy0, fy1 = dfs(y) # 递归处理子节点

    # ----- 情况1:不选 x,则子节点 y 可选可不选 -----
    nf0 = [0] * k
    for i in range(k):
    v = (fy0[i] + fy1[i]) % MOD
    if v == 0:
    continue
    for j in range(k):
    if f0[j] == 0:
    continue
    s = (i + j) % k
    nf0[s] = (nf0[s] + v * f0[j]) % MOD

    # ----- 情况2:选 x,则子节点 y 不能选 -----
    nf1 = [0] * k
    for i in range(k):
    v = fy0[i] # 子节点只能取不选 y 的方案
    if v == 0:
    continue
    for j in range(k):
    if f1[j] == 0:
    continue
    s = (i + j) % k
    nf1[s] = (nf1[s] + v * f1[j]) % MOD

    f0, f1 = nf0, nf1

    return f0, f1

    f0, f1 = dfs(0)
    # 根节点结果 = 不选根 + 选根,再减去空集(1 种)
    return (f0[0] + f1[0] - 1) % MOD

    # 示例测试
    if __name__ == "__main__":
    parent = [-1, 0, 0, 0]
    nums = [2, 1, 2, 1]
    k = 3
    print(countValidSubsets(parent, nums, k))

    2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个

    C++完整代码如下:

      
    
    using namespace std;

    const long long MOD = 1000000007LL;

    pair , vector > dfs( int x, const vector int >>& g, const vector< int >& nums, int k) {
    vector f0(k, 0 ), f1(k, 0 );
    f0[ 0 ] = 1 ; // 不选 x 的空集
    f1[nums[x] % k] = 1 ; // 选 x 的集合(仅包含 x)

    for ( int y : g[x]) {
    auto [fy0, fy1] = dfs(y, g, nums, k);

    // 不选 x,则子节点 y 可选可不选
    vector nf0(k, 0 );
    for ( int i = 0 ; i < k; ++i) {
    long long v = (fy0[i] + fy1[i]) % MOD;
    if (v == 0 ) continue ;
    for ( int j = 0 ; j < k; ++j) {
    if (f0[j] == 0 ) continue ;
    int s = (i + j) % k;
    nf0[s] = (nf0[s] + v * f0[j]) % MOD;
    }
    }

    // 选 x,则子节点 y 不能选
    vector nf1(k, 0 );
    for ( int i = 0 ; i < k; ++i) {
    long long v = fy0[i]; // 只能选 y 中不选 y 的方案
    if (v == 0 ) continue ;
    for ( int j = 0 ; j < k; ++j) {
    if (f1[j] == 0 ) continue ;
    int s = (i + j) % k;
    nf1[s] = (nf1[s] + v * f1[j]) % MOD;
    }
    }

    f0 = move(nf0);
    f1 = move(nf1);
    }

    return {f0, f1};
    }

    int countValidSubsets( const vector< int >& parent, const vector< int >& nums, int k) {
    int n = parent.size();
    vector int >> g(n);
    for ( int i = 1 ; i < n; ++i) {
    int p = parent[i];
    g[p].push_back(i);
    }

    auto [f0, f1] = dfs( 0 , g, nums, k);
    long long ans = (f0[ 0 ] + f1[ 0 ] - 1 ) % MOD; // 减去空集
    if (ans < 0 ) ans += MOD;
    return ( int )ans;
    }

    int main() {
    vector< int > parent = { -1 , 0 , 0 , 0 };
    vector< int > nums = { 2 , 1 , 2 , 1 };
    int k = 3 ;
    int result = countValidSubsets(parent, nums, k);
    cout << result << endl;
    return 0 ;
    }

    2026-08-31:统计有根树中不相邻子集的数目。用go语言,给定一棵包含 n 个节点的有根树,节点编号为 0 到 n-1,其中 0 号节点是根。每个

    我们相信人工智能为普通人提供了一种“增强工具”,并致力于分享全方位的AI知识。在这里,您可以找到最新的AI科普文章、工具评测、提升效率的秘籍以及行业洞察。 欢迎关注“福大大架构师每日一题”,发消息可获得面试资料,让AI助力您的未来发展。

    © 版权声明

    相关文章