1916. 蚁群建房间的方案数

Hard

题目描述

你是一只蚂蚁，需要为你的蚁群添加 n 个编号为 0 到 n-1 的新房间。给定一个长度为 n 的 0 索引整数数组 prevRoom 作为扩建计划，其中 prevRoom[i] 表示在建造房间 i 之前必须先建造房间 prevRoom[i]，并且这两个房间必须直接相连。房间 0 已经建好，所以 prevRoom[0] = -1。扩建计划保证一旦所有房间都建好，从房间 0 可以到达每个房间。

你一次只能建造一个房间，并且只能在已经建好且相连的房间之间自由移动。你可以选择建造任何房间，只要它的前置房间已经建好。

返回建造所有房间的不同顺序数。由于答案可能很大，返回结果对 10^9 + 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 <= 10^5
prevRoom[0] == -1
对于所有 1 <= i < n，0 <= prevRoom[i] < n
一旦所有房间都建好，从房间 0 可以到达每个房间

解题思路

这是一个经典的树形动态规划 + 组合数学问题。

核心思路：

根据 prevRoom 数组构建树结构，房间 0 为根节点
对于每个节点，计算其子树中所有节点的建造方案数
关键洞察：对于一个有多个子树的节点，这些子树的建造顺序可以交叉进行，这就是一个多重排列组合问题

算法步骤：

构建邻接表表示树结构
DFS 计算每个子树的大小和方案数
对于每个节点，使用组合数学计算其所有子树的交叉排列方案数
公式：如果一个节点有多个子树，大小分别为 s1, s2, …, sk，那么这些子树的交叉排列方案数为 C(s1+s2+…+sk, s1) × C(s2+…+sk, s2) × …

组合数计算：

预计算阶乘和阶乘的逆元，用于快速计算组合数
使用费马小定理计算逆元：a^(-1) ≡ a^(p-2) (mod p)

这种方法的时间复杂度为 O(n)，空间复杂度为 O(n)。

代码实现

class Solution {
public:
    const int MOD = 1e9 + 7;
    vector<long long> fact, inv_fact;
    vector<vector<int>> children;
    
    long long power(long long a, long long b) {
        long long res = 1;
        while (b > 0) {
            if (b & 1) res = res * a % MOD;
            a = a * a % MOD;
            b >>= 1;
        }
        return res;
    }
    
    void precompute(int n) {
        fact.resize(n + 1);
        inv_fact.resize(n + 1);
        fact[0] = 1;
        for (int i = 1; i <= n; i++) {
            fact[i] = fact[i - 1] * i % MOD;
        }
        inv_fact[n] = power(fact[n], MOD - 2);
        for (int i = n - 1; i >= 0; i--) {
            inv_fact[i] = inv_fact[i + 1] * (i + 1) % MOD;
        }
    }
    
    long long C(int n, int k) {
        if (k > n || k < 0) return 0;
        return fact[n] * inv_fact[k] % MOD * inv_fact[n - k] % MOD;
    }
    
    pair<int, long long> dfs(int node) {
        int size = 1;
        long long ways = 1;
        
        for (int child : children[node]) {
            auto [child_size, child_ways] = dfs(child);
            ways = ways * child_ways % MOD * C(size + child_size - 1, child_size) % MOD;
            size += child_size;
        }
        
        return {size, ways};
    }
    
    int waysToBuildRooms(vector<int>& prevRoom) {
        int n = prevRoom.size();
        children.resize(n);
        
        for (int i = 1; i < n; i++) {
            children[prevRoom[i]].push_back(i);
        }
        
        precompute(n);
        return dfs(0).second;
    }
};

class Solution:
    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
        MOD = 10**9 + 7
        n = len(prevRoom)
        
        # Build adjacency list
        children = [[] for _ in range(n)]
        for i in range(1, n):
            children[prevRoom[i]].append(i)
        
        # Precompute factorials and inverse factorials
        fact = [1] * (n + 1)
        for i in range(1, n + 1):
            fact[i] = fact[i - 1] * i % MOD
        
        def power(a, b):
            res = 1
            while b:
                if b & 1:
                    res = res * a % MOD
                a = a * a % MOD
                b >>= 1
            return res
        
        inv_fact = [1] * (n + 1)
        inv_fact[n] = power(fact[n], MOD - 2)
        for i in range(n - 1, -1, -1):
            inv_fact[i] = inv_fact[i + 1] * (i + 1) % MOD
        
        def C(n, k):
            if k > n or k < 0:
                return 0
            return fact[n] * inv_fact[k] % MOD * inv_fact[n - k] % MOD
        
        def dfs(node):
            size = 1
            ways = 1
            
            for child in children[node]:
                child_size, child_ways = dfs(child)
                ways = ways * child_ways % MOD * C(size + child_size - 1, child_size) % MOD
                size += child_size
            
            return size, ways
        
        return dfs(0)[1]

public class Solution {
    private const int MOD = 1000000007;
    private long[] fact;
    private long[] invFact;
    private List<int>[] children;
    
    public int WaysToBuildRooms(int[] prevRoom) {
        int n = prevRoom.Length;
        children = new List<int>[n];
        for (int i = 0; i < n; i++) {
            children[i] = new List<int>();
        }
        
        for (int i = 1; i < n; i++) {
            children[prevRoom[i]].Add(i);
        }
        
        Precompute(n);
        return (int)DFS(0).ways;
    }
    
    private void Precompute(int n) {
        fact = new long[n + 1];
        invFact = new long[n + 1];
        fact[0] = 1;
        for (int i = 1; i <= n; i++) {
            fact[i] = fact[i - 1] * i % MOD;
        }
        invFact[n] = Power(fact[n], MOD - 2);
        for (int i = n - 1; i >= 0; i--) {
            invFact[i] = invFact[i + 1] * (i + 1) % MOD;
        }
    }
    
    private long Power(long a, long b) {
        long res = 1;
        while (b > 0) {
            if ((b & 1) == 1) res = res * a % MOD;
            a = a * a % MOD;
            b >>= 1;
        }
        return res;
    }
    
    private long C(int n, int k) {
        if (k > n || k < 0) return 0;
        return fact[n] * invFact[k] % MOD * invFact[n - k] % MOD;
    }
    
    private (int size, long ways) DFS(int node) {
        int size = 1;
        long ways = 1;
        
        foreach (int child in children[node]) {
            var (childSize, childWays) = DFS(child);
            ways = ways * childWays % MOD * C(size + childSize - 1, childSize) % MOD;
            size += childSize;
        }
        
        return (size, ways);
    }
}

var waysToBuildRooms = function(prevRoom) {
    const MOD = 1e9 + 7;
    const n = prevRoom.length;
    
    // Build adjacency list
    const children = Array(n).fill().map(() => []);
    for (let i = 1; i < n; i++) {
        children[prevRoom[i]].push(i);
    }
    
    // Precompute factorials and inverse factorials
    const fact = new Array(n + 1);
    fact[0] = 1;
    for (let i = 1; i <= n; i++) {
        fact[i] = (fact[i - 1] * i) % MOD;
    }
    
    function power(a, b) {
        let res = 1;
        while (b > 0) {
            if (b & 1) res = (res * a) % MOD;
            a = (a * a) % MOD;
            b >>= 1;
        }
        return res;
    }
    
    const invFact = new Array(n + 1);
    invFact[n] = power(fact[n], MOD - 2);
    for (let i = n - 1; i >= 0; i--) {
        invFact[i] = (invFact[i + 1] * (i + 1)) % MOD;
    }
    
    function C(n, k) {
        if (k > n || k < 0) return 0;
        return (fact[n] * invFact[k] % MOD * invFact[n - k]) % MOD;
    }
    
    function dfs(node) {
        let size = 1;
        let ways = 1;
        
        for (const child of children[node]) {
            const [childSize, childWays] = dfs(child);
            ways = (ways * childWays % MOD * C(size + childSize - 1, childSize)) % MOD;
            size += childSize;
        }
        
        return [size, ways];
    }
    
    return dfs(0)[1];
};

在 LeetCode 上运行：力扣中文版 LeetCode 国际版

复杂度分析

复杂度类型	分析
时间复杂度	O(n)
空间复杂度	O(n)

详细分析：

时间复杂度：O(n)，其中 n 是房间数量。预计算阶乘需要 O(n)，DFS 遍历每个节点一次需要 O(n)，每次组合数计算为 O(1)
空间复杂度：O(n)，用于存储树的邻接表、阶乘数组、递归调用栈等

题目描述#

解题思路#

代码实现#

复杂度分析#

相关题目#

题目描述

解题思路

代码实现

复杂度分析

相关题目