【牛客】树的距离 树上主席树

传送门：

题意

$给一颗树，求以x为子树中，距离x大于等于k的点与x的距离和。$

思路

$这题求的是子节点到x的距离，而我们dfs的过程中很容易得到子节点到根的距离dis。$

$所以我们转换一下，求x子节点中距离根大于等于k+dis[x]的点与x的距离和。$

$假设我们知道个数为sum\_num，这些点距离根的和为sum\_dis。$
$则ans=sum\_dis-sum\_num\ast dis[x]$

$怎么求个数呢？怎么求大于等于k的点的个数呢？$

$我们可以利用主席树帮助求解，以dis[x]为叶子的主席树。$
$然后求区间内\ge k的个数即可。$

$这里有几个小细节：$

• $dis的大小会爆int，所以我们需要离散化进而建席树$
• $dfs序建主席树，因为是树上$

$最后，这题需要一定的代码能力。$

Code（709MS）

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

typedef long long ll;

const int N = 2e5 + 10;

struct Edge {
    int v;
    ll w;
};

vector<Edge> g[N];

struct Tree {
    int l, r;
    ll sum, val;
}t[N * 40];
int root[N], cnt;

vector<ll> v;
int getid(ll x) { return lower_bound(v.begin(), v.end(), x) - v.begin() + 1; }

void insert(int pre, int &now, int l, int r, ll p, int val) {
    t[now = ++cnt] = t[pre];
    t[now].val += val;
    t[now].sum += v[p - 1];
    if(l == r) return ;
    int m = (l + r) >> 1;
    if(p <= m) insert(t[pre].l, t[now].l, l, m, p, val);
    else insert(t[pre].r, t[now].r, m + 1, r, p, val);
}

void query(int pre, int now, int ql, int qr, int l, int r, ll &sum_num, ll &sum_dis) {
    if(ql <= l && r <= qr) {
        sum_num += t[now].val - t[pre].val;
        sum_dis += t[now].sum - t[pre].sum;
        return ;
    }
    int m = (l + r) >> 1;
    if(ql <= m) query(t[pre].l, t[now].l, ql, qr, l, m, sum_num, sum_dis);
    if(qr > m)  query(t[pre].r, t[now].r, ql, qr, m + 1, r, sum_num, sum_dis);
}

int in[N], out[N], tim;
ll dis[N];
void dfs(int u, int fa) {
    in[u] = ++tim;
    v.push_back(dis[u]);
    for(auto e : g[u]) {
        int v = e.v;
        if(v == fa) continue;
        dis[v] = dis[u] + e.w;
        dfs(v, u);
    }
    out[u] = tim;
}

void dfs2(int u, int fa, int cct) {
    insert(root[in[u] - 1], root[in[u]], 1, cct, getid(dis[u]), 1);
    for(auto e : g[u]) if(e.v != fa) dfs2(e.v, u, cct);
}

void solve() {
    int n; cin >> n;
    for(int i = 2;i <= n; i++) {
        int p; ll d; cin >> p >> d;
        g[i].push_back(Edge{p, d});
        g[p].push_back(Edge{i, d});
    }
    dfs(1, 0);
    sort(v.begin(), v.end());
    v.erase(unique(v.begin(), v.end()), v.end());
    int cct = (int)v.size();
    dfs2(1, 0, cct);
    int m; cin >> m;
    while(m--) {
        int x; ll y; cin >> x >> y;
        y = dis[x] + y;
        int id = getid(y);
        if(id > cct) {
            cout << 0 << endl;
            continue;
        }
        ll sum_num = 0, sum_dis = 0;
        query(root[in[x] - 1], root[out[x]], id, cct, 1, cct, sum_num, sum_dis);
        cout << (ll)(sum_dis - sum_num * dis[x]) << endl;
    }
}

signed main() {
    solve();
}