统计为蚁群构筑房间的不同顺序

标签: 拓扑排序 数学 动态规划 组合数学

难度: Hard

你是一只蚂蚁,负责为蚁群构筑 n 间编号从 0n-1 的新房间。给你一个 下标从 0 开始 且长度为 n 的整数数组 prevRoom 作为扩建计划。其中,prevRoom[i] 表示在构筑房间 i 之前,你必须先构筑房间 prevRoom[i] ,并且这两个房间必须 直接 相连。房间 0 已经构筑完成,所以 prevRoom[0] = -1 。扩建计划中还有一条硬性要求,在完成所有房间的构筑之后,从房间 0 可以访问到每个房间。

你一次只能构筑 一个 房间。你可以在 已经构筑好的 房间之间自由穿行,只要这些房间是 相连的 。如果房间 prevRoom[i] 已经构筑完成,那么你就可以构筑房间 i

返回你构筑所有房间的 不同顺序的数目 。由于答案可能很大,请返回对 109 + 7 取余 的结果。

示例 1:

输入:prevRoom = [-1,0,1]
输出:1
解释:仅有一种方案可以完成所有房间的构筑:0 → 1 → 2

示例 2:

输入:prevRoom = [-1,0,0,1,2]
输出:6
解释:
有 6 种不同顺序:
0 → 1 → 3 → 2 → 4
0 → 2 → 4 → 1 → 3
0 → 1 → 2 → 3 → 4
0 → 1 → 2 → 4 → 3
0 → 2 → 1 → 3 → 4
0 → 2 → 1 → 4 → 3

提示:

  • n == prevRoom.length
  • 2 <= n <= 105
  • prevRoom[0] == -1
  • 对于所有的 1 <= i < n ,都有 0 <= prevRoom[i] < n
  • 题目保证所有房间都构筑完成后,从房间 0 可以访问到每个房间

Submission

运行时间: 1360 ms

内存: 67.5 MB

MOD = int(1e9 + 7)
N = int(1e5 + 5)
fac, inv = [0] * N, [0] * N
fac[0] = inv[0] = 1
for i in range(1, N):
    fac[i] = fac[i - 1] * i % MOD
    inv[i] = pow(fac[i], MOD - 2, MOD)

class Solution:
    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
        n = len(prevRoom)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[prevRoom[i]].append(i)
        
        f, sz = [1] * n, [1] * n 
        def dfs(u):
            for v in g[u]:
                dfs(v)
                f[u] = f[u] * f[v] * inv[sz[v]] % MOD
                sz[u] += sz[v]
            f[u] = f[u] * fac[sz[u] - 1] % MOD
        
        dfs(0)
        return f[0]

Explain

题解使用了树形 DP 的方法来解决问题。首先将房间之间的依赖关系转换为一棵树的形式,其中节点代表房间,边代表构建依赖。然后从根节点(房间0)开始进行深度优先搜索(DFS),在每个节点上计算构建该节点及其子树的所有不同顺序的数量。这个数量等于该节点所有子节点的顺序数量的乘积,再乘以该节点子树大小的阶乘,并且要考虑排列的顺序,所以还需要除以每个子节点子树大小的阶乘。

时间复杂度: O(n)

空间复杂度: O(n)

MOD = int(1e9 + 7)
N = int(1e5 + 5)
fac, inv = [0] * N, [0] * N
fac[0] = inv[0] = 1
for i in range(1, N):
    fac[i] = fac[i - 1] * i % MOD
    inv[i] = pow(fac[i], MOD - 2, MOD)

class Solution:
    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
        n = len(prevRoom)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[prevRoom[i]].append(i)
        
        f, sz = [1] * n, [1] * n 
        def dfs(u):
            for v in g[u]:
                dfs(v)
                f[u] = f[u] * f[v] * inv[sz[v]] % MOD
                sz[u] += sz[v]
            f[u] = f[u] * fac[sz[u] - 1] % MOD
        
        dfs(0)
        return f[0]

Explore

本题解未特别处理循环依赖的情况,因为题目假设中通常假定输入的依赖关系能够形成一个有效的有向无环图(DAG),即不存在循环依赖。如果实际存在循环依赖,则构建的依赖关系不是有效的树结构,这会导致DFS无法正常执行,可能引发无限递归或其他逻辑错误。在实际应用中,如果需要处理可能包含循环依赖的数据,应该先使用拓扑排序或其他方法验证输入数据是否形成有向无环图。

`sz[u]`数组用于存储以节点`u`为根节点的子树中的节点总数。在DFS过程中,每当访问一个节点`v`,就将`v`的子树大小加到其父节点`u`的子树大小中。`f[u]`数组存储的是以`u`为根节点的子树中所有不同的构建顺序的数量。在计算`f[u]`时,首先将子节点`v`的所有可能的构建顺序(`f[v]`)乘上到目前为止的结果,并考虑不同的排列方式,使用逆元来除以子节点的子树大小的阶乘,以消除重复计数。最后,再乘以当前节点下所有子节点加上自身的阶乘,以考虑所有子节点间的排列组合。

使用`pow(fac[i], MOD - 2, MOD)`计算逆元是基于费马小定理,它指出如果`p`是一个质数且`a`是不被`p`整除的整数,则`a^(p-1) ≡ 1 (mod p)`。由此可以推导出`a^(p-2) ≡ a^(-1) (mod p)`,即`a`的逆元。在此题解中,`MOD`是一个质数(1e9+7),因此可以安全地使用此公式来计算阶乘的模逆,以便进行组合数计算。这种方法简单且在计算组合数时非常有效,特别是在需要大量逆元计算的场景中。