CSES 2081：Fixed-Length Paths II

CSES 2081：Fixed-Length Paths II

題目大意：給你一棵樹，問你總共有幾個路徑使得兩個節點的距離介在 $k_1$ 和 $k_2$ 之間。

解法：重心分治 $+ \text{BIT}$。

$\text{Code：}$

#pragma GCC optimize("O3")
#include <bits/stdc++.h>
using namespace std;

#define LL long long

const int SIZE = 200005;

int n, k1, k2;
LL ans;
vector<int> adj[SIZE];
int mxd = SIZE - 1, cnt[SIZE], wei[SIZE], sum[SIZE];
bool vs[SIZE];

struct BIT {
    void update (int pos, int x) {
        for (; pos <= n; pos += pos & -pos)
            sum[pos] += x;
    }
    int query (int pos) {
        if (pos < 0)
            return 0;
        int re = 0;
        for (;; pos -= pos & -pos) {
            re += sum[pos];
            if (pos == 0)
                return re;
        }
    }
} bit;

int get_siz (int pos, int f) {
    wei[pos] = 1;
    for (int np : adj[pos]) {
        if (np != f && !vs[np])
            wei[pos] += get_siz (np, pos);
    }
    return wei[pos];
}

int get_cent (int pos, int half, int f) {
    for (int np : adj[pos]) {
        if (np != f && !vs[np] && wei[np] >= half)
            return get_cent (np, half, pos);
    }
    return pos;
}

void get_cnt (int pos, int dep, bool cnting, int f) {
    if (dep > k2)
        return;
    mxd = max (mxd, dep);
    if (cnting)
        ans += bit.query (k2 - dep) - bit.query (k1 - dep - 1);
    else
        bit.update (dep, 1);
    for (int np : adj[pos]) {
        if (np != f && !vs[np])
            get_cnt (np, dep + 1, cnting, pos);
    }
}

void Divide (int pos) {
    int cent = get_cent (pos, get_siz (pos, pos) / 2, pos);
    vs[cent] = 1;
    mxd = 0;
    for (int np : adj[cent]) {
        if (!vs[np]) {
            get_cnt (np, 1, 1, np);
            get_cnt (np, 1, 0, np);
        }
    }
    for (int i = 1; i <= mxd; i++)
        bit.update (i, -bit.query (i) + 1);
    for (int np : adj[cent])
        if (!vs[np])
            Divide (np);
}

void solve() {
    cin >> n >> k1 >> k2;
    sum[0] = 1;
    for (int i = 1; i < n; i++) {
        int a, b;
        cin >> a >> b;
        adj[a].push_back (b);
        adj[b].push_back (a);
    }
    Divide (1);
    cout << ans << '\n';
}

int main() {
    ios::sync_with_stdio (false), cin.tie (0);
    solve();
}

我的分享就到這裡結束了，如果喜歡我的 $\text{Blog}$，歡迎追蹤！