740-D. Alyona and a tree

2021-03-18 21:40:57 +0800 CST Tags: binary search, difference-array, prefix-sum, tree

出處：https://codeforces.com/contest/740/problem/D
提交：https://codeforces.com/contest/740/submission/110308918

問題：給定 $n\in[1,2\cdot 10^5]$ 個頂點 (1 ~ n) 以及它們各自儲存一個介於 $[1,10^9]$ 的數值，再接下來從第 2 個開始到最後的頂點，都會有自己的父/母頂點以及連接兩者的邊權重、也介於 $[1,10^9]$，第一個頂點是為樹根所以不用給資訊，而且保證輸入的圖是一棵樹。現在定義頂點 u 可以「控制」頂點 v 的條件為 v 在 u 底下 (兩者不相等) 而且 u 走到 v 的邊權重和不超過 v 儲存的數值。請印出每個頂點各自可以控制的頂點數量。

解法：首先可以觀察到愈靠近奴隸 v 的祖先愈有權力可以控制 v，所以對於每個非根頂點的奴隸 v，我都可以在從它往上走到根頂點的這段路上進行 binary search，找到最高可以控制它的主人，此時可以獲得這區間所有合法的主人。一個直觀的作法是直接把數量更新進這些主人的計數器裡，只是如此一來對每個奴隸都要更新這麼多人的資訊，會達到 $\mathcal O(n^2)$ 的時間複雜度，因此這邊我們是採用 difference array 技巧，這個技巧真的好用，我在這個奴隸 v 的最靠近合法主人計數器設下 +1，意味著由下往上累加每個主人所擁有的奴隸數量時，從這個奴隸 v 開始都要考慮到它，之後在 v 的最遠離合法主人的「父/母頂點」計數器設下 -1，意味著這個奴隸 v 已經不再受到從現在開始往上走的途中所遇到的所有頂點的控制了。這邊在思考的時候建議把 +1 和 -1 想成是多鋪一條繩子，所以計數器可以看成是有多少條繩子疊在上面，這樣子情感上比較可以相信程式碼中 dfs2 的累加，換句話說 difference array 也能用在樹上！實作的時候只要注意一下邊界條件，例如根頂點因為不受控制所以不用參與搜尋，以及可能有頂點完全不受任何一個頂點控制，還有 binary search 要如何縮小解區間比較方便等等，就可以了。

實作：

 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


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

int n, a[200001], w[200001], dp[200001];
vector<int> adj[200001], indices;
vector<long long> sum;

void dfs1(int v) {
    indices.push_back(v); sum.push_back((sum.size()==0 ? 0 : sum.back()) + w[v]);
    /* use binary search to find the smallest index "hi"
     * such that sum.back() - sum[hi] <= a[v]. Note that
     * such index must exist because a[v] >= 1 and the LHS
     * can be 0 if hi == sum.size()-1.
     */
    if (v > 1) { // the root vertex cannot be controlled.
        int lo = 0, hi = indices.size()-1;
        while (lo < hi) {
            int m = (lo + hi) / 2;
            if (sum.back() - sum[m] <= a[v])
                hi = m;
            else
                lo = m + 1;
        }
        dp[indices[indices.size()-2]] += 1;
        if (hi >= 1) dp[indices[hi-1]] -= 1;
    }
    for (auto u: adj[v])
        dfs1(u);
    indices.pop_back(); sum.pop_back();
}

int dfs2(int v) {
    for (auto u: adj[v])
        dp[v] += dfs2(u);
    return dp[v];
}

int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); // IO 優化
    cin >> n;
    for (int i=1; i<=n; i++) cin >> a[i];
    for (int i=2; i<=n; i++) {
        cin >> w[i]; // parent
        adj[w[i]].push_back(i); // push child
        cin >> w[i];
    }
    dfs1(1); dfs2(1);
    for (int i=1; i<=n; i++) cout << dp[i] << " ";
    return 0;
}

結果： Example image