来源第 247 场周赛 Q4难度困难分数2486
题目描述
你是一只蚂蚁，负责为蚁群构筑 n 间编号从 0 到 n-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.length2 <= n <= 105prevRoom[0] == -1- 对于所有的
1 <= i < n ，都有 0 <= prevRoom[i] < n - 题目保证所有房间都构筑完成后，从房间
0 可以访问到每个房间
解法
方法一
思考
建造顺序必须尊重树形先序约束，合法序列是该有根树的拓扑序。枚举排列再检验在 \(n\le 10^5\) 时不可行。
子树内部的建造顺序相互独立，合并两棵子树时只需在合并后的位置中为其中一棵选择下标，方案数为组合数。多子树则从左到右陆续合并。
DFS 返回子树大小，对每个结点用 \(\binom{s+t}{t}\) 把已合并大小 \(s\) 与新子树大小 \(t\) 相乘，并对模数取余。
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 | class Solution:
def waysToBuildRooms(self, prevRoom: List[int]) -> int:
modulo = 10**9 + 7
ingoing = defaultdict(set)
outgoing = defaultdict(set)
for i in range(1, len(prevRoom)):
ingoing[i].add(prevRoom[i])
outgoing[prevRoom[i]].add(i)
ans = [1]
def recurse(i):
if len(outgoing[i]) == 0:
return 1
nodes_in_tree = 0
for v in outgoing[i]:
cn = recurse(v)
if nodes_in_tree != 0:
ans[0] *= comb(nodes_in_tree + cn, cn)
ans[0] %= modulo
nodes_in_tree += cn
return nodes_in_tree + 1
recurse(0)
return ans[0] % modulo
|
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 | class Solution {
private static final int MOD = 1_000_000_007;
private List<Integer>[] g;
private long[] fact;
private long[] invFact;
private long ans = 1;
public int waysToBuildRooms(int[] prevRoom) {
int n = prevRoom.length;
g = new List[n];
Arrays.setAll(g, i -> new ArrayList<>());
for (int i = 1; i < n; ++i) {
g[prevRoom[i]].add(i);
}
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] = qpow(fact[n], MOD - 2);
for (int i = n; i > 0; --i) {
invFact[i - 1] = invFact[i] * i % MOD;
}
dfs(0);
return (int) ans;
}
private int dfs(int u) {
int merged = 0;
for (int v : g[u]) {
int cn = dfs(v);
if (merged != 0) {
ans = ans * comb(merged + cn, cn) % MOD;
}
merged += cn;
}
return merged + 1;
}
private long comb(int n, int k) {
return fact[n] * invFact[k] % MOD * invFact[n - k] % MOD;
}
private long qpow(long a, long n) {
long ans = 1;
for (; n > 0; n >>= 1) {
if ((n & 1) == 1) {
ans = ans * a % MOD;
}
a = a * a % MOD;
}
return ans;
}
}
|
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 | class Solution {
public:
int waysToBuildRooms(vector<int>& prevRoom) {
int n = prevRoom.size();
g.assign(n, {});
for (int i = 1; i < n; ++i) {
g[prevRoom[i]].push_back(i);
}
fact.assign(n + 1, 1);
invFact.assign(n + 1, 1);
for (int i = 1; i <= n; ++i) {
fact[i] = fact[i - 1] * i % MOD;
}
invFact[n] = qpow(fact[n], MOD - 2);
for (int i = n; i > 0; --i) {
invFact[i - 1] = invFact[i] * i % MOD;
}
ans = 1;
dfs(0);
return ans;
}
private:
static constexpr int MOD = 1e9 + 7;
vector<vector<int>> g;
vector<long long> fact, invFact;
int ans;
int dfs(int u) {
int merged = 0;
for (int v : g[u]) {
int cn = dfs(v);
if (merged) {
ans = 1LL * ans * comb(merged + cn, cn) % MOD;
}
merged += cn;
}
return merged + 1;
}
long long comb(int n, int k) {
return fact[n] * invFact[k] % MOD * invFact[n - k] % MOD;
}
long long qpow(long long a, long long n) {
long long res = 1;
for (; n; n >>= 1) {
if (n & 1) {
res = res * a % MOD;
}
a = a * a % MOD;
}
return res;
}
};
|
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 | func waysToBuildRooms(prevRoom []int) int {
const mod = 1_000_000_007
n := len(prevRoom)
g := make([][]int, n)
for i := 1; i < n; i++ {
g[prevRoom[i]] = append(g[prevRoom[i]], i)
}
fact := make([]int, n+1)
invFact := make([]int, n+1)
fact[0] = 1
for i := 1; i <= n; i++ {
fact[i] = fact[i-1] * i % mod
}
qpow := func(a, n int) int {
res := 1
for ; n > 0; n >>= 1 {
if n&1 == 1 {
res = res * a % mod
}
a = a * a % mod
}
return res
}
invFact[n] = qpow(fact[n], mod-2)
for i := n; i > 0; i-- {
invFact[i-1] = invFact[i] * i % mod
}
comb := func(n, k int) int {
return fact[n] * invFact[k] % mod * invFact[n-k] % mod
}
ans := 1
var dfs func(int) int
dfs = func(u int) int {
merged := 0
for _, v := range g[u] {
cn := dfs(v)
if merged != 0 {
ans = ans * comb(merged+cn, cn) % mod
}
merged += cn
}
return merged + 1
}
dfs(0)
return ans
}
|