您现在的位置：程式碼之家 >> 編程語言 > >> C語言 >> C++ >> C++ 入門

BZOJ 2783 [JLOI2012] 樹：路徑和為 S 的路徑計數

編輯：C++ 入門

問題描述

給定一棵 N 個節點的有根樹（根為 1，深度 0），每個節點有一個正整數權值。求節點權值和恰好為 S 的路徑數量，路徑上節點深度必須嚴格遞增（即路徑自祖先向後代方向）。N ≤ 100,000，權值與 S 均 ≤ 1000。

核心觀察與前綴和

對於任意節點 u，從根到 u 的路徑上權值前綴和記為 sum[u]。若一條路徑從祖先 v 到後代 u（v 為 u 的祖先，含 v=u），其權值和為 sum[u] - sum[parent[v]]（約定 sum[0]=0）。因此問題轉化為：對每個節點 u，統計有多少祖先 v 滿足 sum[u] - sum[parent[v]] == S，即 sum[parent[v]] == sum[u] - S。

倍增 LCA 快速跳轉祖先

由於 N 較大，逐個向上遍歷祖先會超時。預處理倍增表 f[u][k] 表示 u 的第 2^k 個祖先。查詢 u 的第 k 個祖先時，將 k 拆解為二進制，按位跳躍，時間 O(log N)。

二分查找目標祖先

對於固定 u，sum[parent[v]] 隨著 v 從根向下移動單調遞增（權值為正）。因此可在深度範圍 [1, depth[u]] 上二分查找滿足 sum[parent[v]] == sum[u] - S 的祖先 v。結合倍增跳轉，單次查詢 O(log² N)，總複雜度 O(N log² N)，可通過。

代碼框架（C++）

#include <bits/stdc++.h>
using namespace std;
const int N = 100005, LOG = 18;
int n, S, val[N], f[N][LOG], depth[N], sum[N];
vector g[N];
int ans = 0;

void dfs(int u, int fa) {
    f[u][0] = fa;
    depth[u] = depth[fa] + 1;
    sum[u] = sum[fa] + val[u];
    for (int v : g[u]) if (v != fa) dfs(v, u);
}

int kth_ancestor(int u, int k) {
    for (int i = 0; k; ++i, k >>= 1)
        if (k & 1) u = f[u][i];
    return u;
}

int find_ancestor(int u, int target) {
    int l = 1, r = depth[u] - 1;
    while (l <= r) {
        int mid = (l + r) >> 1;
        int anc = kth_ancestor(u, mid);
        if (sum[anc] >= target) r = mid - 1;
        else l = mid + 1;
    }
    int anc = kth_ancestor(u, l);
    return (sum[anc] == target) ? anc : 0;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n >> S;
    for (int i = 1; i <= n; ++i) cin >> val[i];
    for (int i = 1; i < n; ++i) {
        int x, y; cin >> x >> y;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    dfs(1, 0);
    for (int j = 1; j < LOG; ++j)
        for (int i = 1; i <= n; ++i)
            f[i][j] = f[f[i][j-1]][j-1];
    for (int u = 1; u <= n; ++u) {
        if (val[u] == S) ++ans;
        if (u == 1) continue;
        int target = sum[u] - S;
        if (target < 0) continue;
        int anc = find_ancestor(u, target);
        if (anc) ++ans;
    }
    cout << ans << '\n';
    return 0;
}

優化提示

由於 S ≤ 1000，權值 ≤ 1000，路徑長度實際上不會超過 S（全 1 時長度 S），可將二分上界限制為 min(depth[u]-1, S)，進一步降低常數。
也可用「深度為鍵的哈希表/陣列」維護當前路徑上的前綴和計數，DFS 進出時更新，單次查詢 O(1)，總複雜度 O(N)，但需注意回溯時恢復狀態。