Count Ways to Build Rooms in an Ant Colony (LeetCode 1916)

Nguồn: LeetCode 1916 Độ khó: Khó (Hard) QHĐ trên cây Tổ hợp Nghịch đảo modulo

Đề bài

🇻🇳 Tiếng Việt🇬🇧 English (đề gốc)

1916. Count Ways to Build Rooms in an Ant Colony – Đếm số cách xây các phòng trong tổ kiến (LeetCode, độ khó Hard)

Bạn là một chú kiến được giao nhiệm vụ xây thêm $n$ phòng mới, đánh số từ $0$ đến $n - 1$, cho tổ của mình. Bạn được cho bản kế hoạch mở rộng dưới dạng một mảng số nguyên prevRoom độ dài $n$ (chỉ số tính từ 0), trong đó prevRoom[i] cho biết bạn phải xây phòng prevRoom[i] trước khi xây phòng $i$, và hai phòng này phải được nối trực tiếp với nhau. Phòng $0$ đã được xây sẵn, nên prevRoom[0] = -1. Bản kế hoạch đảm bảo rằng khi xây xong tất cả các phòng thì mọi phòng đều đi tới được từ phòng $0$.

Mỗi lúc bạn chỉ xây được một phòng, và bạn chỉ có thể đi lại tự do giữa các phòng đã xây nếu chúng được nối với nhau. Bạn có thể chọn xây bất kỳ phòng nào miễn là phòng trước nó đã được xây.

Hãy trả về số thứ tự khác nhau để xây tất cả các phòng. Vì đáp án có thể rất lớn, hãy trả về phần dư khi chia cho $10^9 + 7$.

Ràng buộc:

• $n =$ prevRoom.length
• $2 \le n \le 10^5$
• prevRoom[0] == -1
• $0 \le$ prevRoom[i] $\lt n$ với mọi $1 \le i \lt n$
• Khi xây xong tất cả, mọi phòng đều đi tới được từ phòng $0$.

Định dạng vào/ra để chạy thử trên máy: trên LeetCode các bạn chỉ cần viết hàm waysToBuildRooms trong class Solution. Để chạy như một chương trình bình thường (và dùng khung chạy tay bên dưới), mình đặt thêm định dạng:

• Dữ liệu vào: dòng đầu là số nguyên $n$, dòng thứ hai gồm $n$ số nguyên prevRoom[0], prevRoom[1], ..., prevRoom[n-1] (số đầu tiên luôn là $-1$).
• Kết quả: một số nguyên là đáp án của bài.

Ví dụ 1

Input: prevRoom = [-1,0,1]
Output: 1
---stdin---
3
-1 0 1

Chỉ có đúng một cách xây các phòng: 0 → 1 → 2.

Ví dụ 2

Input: prevRoom = [-1,0,0,1,2]
Output: 6
---stdin---
5
-1 0 0 1 2

6 cách là:
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

1916. Count Ways to Build Rooms in an Ant Colony (LeetCode, Hard)

You are an ant tasked with adding n new rooms numbered 0 to n-1 to your colony. You are given the expansion plan as a 0-indexed integer array of length n, prevRoom, where prevRoom[i] indicates that you must build room prevRoom[i] before building room i, and these two rooms must be connected directly. Room 0 is already built, so prevRoom[0] = -1. The expansion plan is given such that once all the rooms are built, every room will be reachable from room 0.

You can only build one room at a time, and you can travel freely between rooms you have already built only if they are connected. You can choose to build any room as long as its previous room is already built.

Return the number of different orders you can build all the rooms in. Since the answer may be large, return it modulo 10^9 + 7.

Example 1:

Input: prevRoom = [-1,0,1]
Output: 1
Explanation: There is only one way to build the additional rooms: 0 → 1 → 2

Example 2:

Input: prevRoom = [-1,0,0,1,2]
Output: 6
Explanation:
The 6 ways are:
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

Constraints:

• n == prevRoom.length
• 2 <= n <= 10^5
• prevRoom[0] == -1
• 0 <= prevRoom[i] < n for all 1 <= i < n
• Every room is reachable from room 0 once all the rooms are built.

Lời mở đầu

Bài này mình được một bạn học viên hỏi nên sẵn tiện mình viết bài giải chi tiết để nhiều bạn được tham khảo luôn. Hy vọng mang đến được giá trị cho các bạn 😇.

Link nộp bài

Bấm vào nút bên dưới để nộp bài và được hệ thống CLUE OJ chấm tự động.

Nộp bài trên CLUE OJ

Lưu ý nhỏ: khi nộp trên CLUE OJ các bạn dán toàn bộ code, gồm cả hàm main, vì đề trên CLUE OJ đọc dữ liệu vào và in kết quả đúng theo định dạng ở phần đề bài, giống như khi các bạn chạy thử trên máy hoặc ở trang Chạy code online của blog.

Hiểu đề cho đúng: vì mỗi phòng $i \ne 0$ có đúng một "phòng trước" prevRoom[i] và mọi phòng đều nối về phòng 0, nên các phòng tạo thành một cây có gốc là phòng 0, cha của $i$ là prevRoom[i]. Một thứ tự xây hợp lệ là một cách xếp các đỉnh sao cho cha luôn đứng trước con. Lưu ý prevRoom[i] có thể lớn hơn $i$ (ví dụ [-1, 2, 0]: phòng 2 phải xây trước phòng 1).

Với ví dụ 2, cây như sau: phòng 0 có hai con là 1 và 2; phòng 3 là con của 1, phòng 4 là con của 2. Nhánh trái phải xây theo thứ tự $1 \to 3$, nhánh phải theo thứ tự $2 \to 4$, nhưng ta được phép xen kẽ hai nhánh với nhau – chính 6 cách trong đề là 6 cách xen kẽ đó.

Cây của ví dụ 2 (prevRoom = [-1,0,0,1,2]).

Lộ trình lời giải, bấm vào từng phần để nhảy tới:

Cách 1Quay lui thử mọi thứ tự xây · O(n!) – chỉ n nhỏCách 2QHĐ trên cây + tổ hợp · O(n·log MOD)

Cách 1Quay lui thử mọi thứ tự xây – O(n!)

Làm theo đúng đề: đánh dấu built[0] = true. Ở mỗi bước, thử xây mọi phòng $i$ chưa xây mà phòng trước nó (prevRoom[i]) đã xây rồi, đệ quy tiếp, quay lui thì bỏ đánh dấu. Khi đã xây đủ $n$ phòng thì đếm thêm 1 cách.

Source code để các bạn tham khảo nếu làm theo hướng này:

#include <bits/stdc++.h>
using namespace std;

const long long MOD = 1000000007;
int n, cnt;
vector<int> par;
bool built[20];

void backtrack(int k) {
    if (k == n) {
        cnt++;
        return;
    }
    for (int i = 1; i < n; i++) {
        if (!built[i] && built[par[i]]) {
            built[i] = true;
            backtrack(k + 1);
            built[i] = false;
        }
    }
}

int main() {
    cin >> n;
    par.resize(n);
    for (int i = 0; i < n; i++) cin >> par[i];
    built[0] = true;
    backtrack(1);
    cout << cnt % MOD;
    return 0;
}

import sys

MOD = 1000000007


def main():
    data = sys.stdin.read().split()
    n = int(data[0])
    par = list(map(int, data[1:1 + n]))
    built = [False] * n
    built[0] = True
    cnt = 0

    def backtrack(k):
        nonlocal cnt
        if k == n:
            cnt += 1
            return
        for i in range(1, n):
            if not built[i] and built[par[i]]:
                built[i] = True
                backtrack(k + 1)
                built[i] = False

    backtrack(1)
    print(cnt % MOD)


if __name__ == "__main__":
    main()

Chỉ chạy được với n rất nhỏ

Số thứ tự xây có thể lên tới $(n - 1)!$ (khi cây hình sao: phòng 0 nối thẳng tới tất cả). Với $n = 12$ đã là gần 40 triệu cách, $n = 10^5$ thì không tưởng. Cách này chỉ dùng để hiểu đề và làm code đối chiếu.

Không gian$O(n)$

Thời gian$O(n \cdot n!)$

Cách 2Quy hoạch động trên cây + tổ hợp – O(n·log MOD)

Phương pháp giải: Quy hoạch động trên cây + Tổ hợp.

Gọi $f[u]$ là số cách xây (số lộ trình hợp lệ) các phòng trong cây con gốc $u$, và $sz[u]$ là số đỉnh của cây con đó. Xét các cây con gốc $v$ (con của $u$): số cách xây riêng trong cây con gốc $v$ là $f[v]$.

Tuy nhiên, đề nói rằng chúng ta không nhất thiết phải xây xong cây con này rồi mới sang cây con khác: đang xây ở cây con này, ta có thể nhảy sang cây con khác. Như vậy, giả sử với cây con gốc $v_1$ ta đã có một lộ trình, với cây con gốc $v_2$ ta có một lộ trình khác, rồi $v_3, v_4, \ldots$ – ta cần đếm số cách trộn tất cả các lộ trình này lại với nhau sao cho vẫn giữ nguyên thứ tự xuất hiện của các đỉnh trong từng lộ trình.

Số cách trộn hai lộ trình

Giả sử ta đã có lộ trình $X$ gồm $sz(X)$ đỉnh, giờ thêm lộ trình $Y$ (của cây con gốc $v$) gồm $sz(Y)$ đỉnh. Số cách trộn $Y$ vào $X$ để được một lộ trình gồm $sz(X) + sz(Y)$ đỉnh mà thứ tự bên trong $X$ và bên trong $Y$ không đổi là:

C(sz(X) + sz(Y), sz(Y))  =  tổ hợp chập sz(Y) của sz(X) + sz(Y)

Ví dụ: lộ trình $X$ có 3 đỉnh X1 → X2 → X3, lộ trình $Y$ có 5 đỉnh Y1 → Y2 → Y3 → Y4 → Y5. Số cách trộn là $C(8, 3) = 56$. Vài lộ trình thỏa mãn:

• Y1 → X1 → X2 → Y2 → X3 → Y3 → Y4 → Y5
• X1 → X2 → X3 → Y1 → Y2 → Y3 → Y4 → Y5
• Y1 → Y2 → Y3 → X1 → X2 → X3 → Y4 → Y5

Chứng minh công thức: trộn 2 lộ trình lại thì có tổng cộng $sz(X) + sz(Y)$ đỉnh, ta xem như có $sz(X) + sz(Y)$ dấu ? xếp hàng: ? → ? → ? → … → ?. Một cách trộn thỏa mãn chính là một cách chọn ra $sz(Y)$ dấu ? để đặt lần lượt các đỉnh của $Y$ vào (theo đúng thứ tự của $Y$), các dấu còn lại đặt lần lượt các đỉnh của $X$. Số cách chọn là $C(sz(X) + sz(Y), sz(Y))$. ∎

Công thức quy hoạch động: duyệt lần lượt từng con $v$ của $u$, cộng dồn $sz[u] \mathrel{+}= sz[v]$ (lúc này $sz[u]$ là tổng số đỉnh các cây con đã trộn, kể cả $v$), rồi:

f[u] = f[u] · f[v] · C(sz[u], sz[v])  (mod 10⁹ + 7)

Trộn xong hết các con thì tăng $sz[u]$ thêm 1 cho chính đỉnh $u$ – đỉnh $u$ luôn phải xây đầu tiên trong cây con của nó nên không có thêm lựa chọn nào. Khởi tạo $f[u] = 1$, đáp án là $f[0]$.

Kiểm tra ví dụ 2: lá 3 và 4 có $f = 1, sz = 1$. Đỉnh 1: $f = 1 \cdot 1 \cdot C(1, 1) = 1$, $sz = 2$; tương tự đỉnh 2: $f = 1$, $sz = 2$. Đỉnh 0: trộn con 1 được $f = 1 \cdot C(2, 2) = 1$, trộn con 2 được $f = 1 \cdot 1 \cdot C(4, 2) = 6$. Đáp án 6 ✔.

Tính C(n, k) theo modulo

Mình tính sẵn $fact[i] = i! \bmod M$ và $invFact[i] = (i!)^{-1} \bmod M$. Vì $M = 10^9 + 7$ là số nguyên tố nên theo định lý Fermat nhỏ, nghịch đảo của $x$ là $x^{M - 2} \bmod M$, tính bằng lũy thừa nhanh POW. Khi đó $C(n, k) = fact[n] \cdot invFact[n - k] \cdot invFact[k] \bmod M$.

Các bạn nhập mảng prevRoom rồi bấm Tiếp để xem từng bước trộn (đỉnh vàng là đỉnh đang xử lý, đỉnh xanh là đã tính xong; dưới mỗi đỉnh ghi f và sz):

prevRoom

Ví dụ nhanh:

Code dưới đây của mình đã AC với tư tưởng đã trình bày ở trên, anh em có thể tham khảo nhé:

#include <bits/stdc++.h>
using namespace std;

const long long MOD = 1000000007;
int sz[100005];
vector<int> child[100005];
long long dp[100005];
long long fact[100005], invFact[100005];
long long POW(long long x, long long n) {
    if (n == 0) {
        return 1LL;
    }
    long long tmp = POW(x, n / 2);
    (tmp *= tmp) %= MOD;
    if (n % 2 == 1) {
        (tmp *= x) %= MOD;
    }
    return tmp;
}
long long C(int n, int k) {
    long long ret = fact[n];
    (ret *= invFact[n - k]) %= MOD;
    (ret *= invFact[k]) %= MOD;
    return ret;
}
void init(int n) {
    fact[0] = invFact[0] = 1LL;
    for (int i = 1; i <= n; i++) {
        fact[i] = (fact[i - 1] * (long long) i) % MOD;
        invFact[i] = POW(fact[i], MOD - 2);
    }
}
void dfs(int u) {
    dp[u] = 1;
    for (int i = 0; i < (int) child[u].size(); i++) {
        int v = child[u][i];
        dfs(v);
        sz[u] += sz[v];
        (dp[u] *= dp[v]) %= MOD;
        (dp[u] *= C(sz[u], sz[v])) %= MOD;
    }
    sz[u]++;
}
class Solution {
public:
    int waysToBuildRooms(vector<int>& prevRoom) {
        int n = (int) prevRoom.size();
        init(n);
        memset(dp, 0, sizeof(dp));
        memset(sz, 0, sizeof(sz));
        for (int i = 0; i < n; i++) {
            child[i].clear();
        }
        for (int i = 1; i < n; i++) {
            child[prevRoom[i]].push_back(i);
        }
        dfs(0);
        return (int) dp[0];
    }
};

int main() {
    int n;
    cin >> n;
    vector<int> prevRoom(n);
    for (int i = 0; i < n; i++) cin >> prevRoom[i];
    Solution sol;
    cout << sol.waysToBuildRooms(prevRoom);
    return 0;
}

import sys
from typing import List

MOD = 1000000007


class Solution:
    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
        n = len(prevRoom)
        fact = [1] * (n + 1)
        for i in range(1, n + 1):
            fact[i] = fact[i - 1] * i % MOD
        inv_fact = [1] * (n + 1)
        inv_fact[n] = pow(fact[n], MOD - 2, MOD)
        for i in range(n, 0, -1):
            inv_fact[i - 1] = inv_fact[i] * i % MOD

        def C(a, b):
            return fact[a] * inv_fact[a - b] % MOD * inv_fact[b] % MOD

        child = [[] for _ in range(n)]
        for i in range(1, n):
            child[prevRoom[i]].append(i)
        # Thứ tự DFS không đệ quy: cha luôn đứng trước con,
        # duyệt ngược lại thì con được tính xong trước cha.
        order = [0]
        for u in order:
            order.extend(child[u])
        sz = [0] * n
        dp = [1] * n
        for u in reversed(order):
            s = 0
            for v in child[u]:
                s += sz[v]
                dp[u] = dp[u] * dp[v] % MOD * C(s, sz[v]) % MOD
            sz[u] = s + 1
        return dp[0]


def main():
    data = sys.stdin.read().split()
    n = int(data[0])
    prevRoom = list(map(int, data[1:1 + n]))
    print(Solution().waysToBuildRooms(prevRoom))


if __name__ == "__main__":
    main()

Giải thích một số chỗ trong code:

1. init(n) tính sẵn $fact[i]$ và $invFact[i] = $ POW(fact[i], MOD - 2) cho $i = 0..n$. Mỗi lần gọi POW tốn khoảng 30 bước nên tổng cộng $O(n \log M)$.
2. Trong dfs(u): sz[u] += sz[v] được làm trước khi nhân C(sz[u], sz[v]), nên lúc đó sz[u] chính là $sz(X) + sz(Y)$ như công thức. Dòng sz[u]++ đặt sau vòng lặp để cộng chính đỉnh $u$.
3. child[prevRoom[i]].push_back(i): dựng danh sách con. Vì prevRoom[i] có thể lớn hơn $i$ nên phải DFS từ gốc chứ không tính theo thứ tự chỉ số được.
4. Bản Python không dùng đệ quy (cây có thể là một đường thẳng sâu $10^5$ tầng, Python sẽ tràn đệ quy). Mình lấy thứ tự duyệt theo chiều rộng từ gốc (cha luôn đứng trước con) rồi xử lý ngược thứ tự đó, nên con luôn được tính xong trước cha. $invFact$ cũng được tính ngược từ $invFact[n]$ cho nhanh.

Khung chạy tay (để dễ theo dõi, giới hạn $n \le 8$; bước tính POW được gom thành một bước):

C++▶ Chạy tay từng dòng codeCách 2: init giai thừa + nghịch đảo, dựng cây con, dfs trộn các cây con bằng C(sz[u], sz[v]).

Dữ liệu vào (stdin) – sửa thoải mái5 -1 0 0 1 2

Dữ liệu mẫu:

1#include <bits/stdc++.h>

2using namespace std;

3

4const long long MOD = 1000000007;

5int sz[100005];

6vector<int> child[100005];

7long long dp[100005];

8long long fact[100005], invFact[100005];

9long long POW(long long x, long long n) {

10 if (n == 0) {

11 return 1LL;

12 }

13 long long tmp = POW(x, n / 2);

14 (tmp *= tmp) %= MOD;

15 if (n % 2 == 1) {

16 (tmp *= x) %= MOD;

17 }

18 return tmp;

19}

20long long C(int n, int k) {

21 long long ret = fact[n];

22 (ret *= invFact[n - k]) %= MOD;

23 (ret *= invFact[k]) %= MOD;

24 return ret;

25}

26void init(int n) {

27 fact[0] = invFact[0] = 1LL;

28 for (int i = 1; i <= n; i++) {

29 fact[i] = (fact[i - 1] * (long long) i) % MOD;

30 invFact[i] = POW(fact[i], MOD - 2);

31 }

32}

33void dfs(int u) {

34 dp[u] = 1;

35 for (int i = 0; i < (int) child[u].size(); i++) {

36 int v = child[u][i];

37 dfs(v);

38 sz[u] += sz[v];

39 (dp[u] *= dp[v]) %= MOD;

40 (dp[u] *= C(sz[u], sz[v])) %= MOD;

41 }

42 sz[u]++;

43}

44class Solution {

45public:

46 int waysToBuildRooms(vector<int>& prevRoom) {

47 int n = (int) prevRoom.size();

48 init(n);

49 memset(dp, 0, sizeof(dp));

50 memset(sz, 0, sizeof(sz));

51 for (int i = 0; i < n; i++) {

52 child[i].clear();

53 }

54 for (int i = 1; i < n; i++) {

55 child[prevRoom[i]].push_back(i);

56 }

57 dfs(0);

58 return (int) dp[0];

59 }

60};

61

62int main() {

63 int n;

64 cin >> n;

65 vector<int> prevRoom(n);

66 for (int i = 0; i < n; i++) cin >> prevRoom[i];

67 Solution sol;

68 cout << sol.waysToBuildRooms(prevRoom);

69 return 0;

70}

ChậmVừaNhanhRất nhanh

Đánh giá độ phức tạp của cách làm này:

• Không gian: các mảng fact, invFact, dp, sz, danh sách con – $O(n)$.
• Thời gian: init gọi POW $n$ lần, mỗi lần $O(\log M)$; DFS đi qua mỗi cạnh một lần, mỗi lần tính C trong $O(1)$ ⇒ tổng $O(n \log M)$, thực tế coi như $O(n)$.

Không gian$O(n)$

Thời gian$O(n \log M)$

Đệ quy sâu

Cây có thể là một đường thẳng $0 \to 1 \to \ldots \to n - 1$, khi đó dfs đệ quy sâu $10^5$ tầng. Code C++ trên LeetCode vẫn chạy ổn (mình thử trên máy với stack 8 MB cũng chạy ổn), nhưng với Python bắt buộc phải khử đệ quy như bản Python ở trên.

Ghi chú của mình: Code C++ và Python của cả 2 cách đã được kiểm tra với 2 ví dụ của đề (đều đúng) và 400 cây ngẫu nhiên $n \le 9$ (cây ngẫu nhiên, đường thẳng, hình sao, cây sâu; nhãn phòng xáo trộn nên có cả prevRoom[i] > i), so với code quay lui Cách 1 – khớp hết. Với $n = 10^5$ (đường thẳng, hình sao, ngẫu nhiên), Cách 2 C++ chạy dưới 0.1 giây và Python cho cùng kết quả.

Rất vui khi anh em nào đã đọc đến những dòng cuối cùng này. Anh em đọc chắc cũng mỏi mắt luôn chứ 😅 vậy anh em nghĩ người viết ra hết đống ở trên còn thế nào nữa 🤪. Vậy nên đừng tiếc gì cho mình một lời bình luận hoặc chia sẻ bài viết để động viên mình nè. Chúc anh em học tốt và gặt hái nhiều thành công 😇