P3899 [湖南集训]谈笑风生

P3899 [湖南集训]谈笑风生
这题定义诡异，似乎对于一对 (x, y) 如果 (y) 是 (x) 在范围 (k) 以内的儿子，那么 (x) 既和它彼此彼此，又高明到不知道哪里去了。
然后你思考一下，分两部分：

1. (x) 是 (y) 的祖先。
2. (y) 是 (x) 的祖先。

容易发现，当你在 (x) 的时候.
第一部分 (sum_1 = min(dep_x - 1, k) cdot (siz_x-1))。

第二部分 (sum_2 = displaystyle sum_{v in subtree_u & dep_v - dep_u le k} siz_v)。

第二部分显然很线段树合并，第一部分 (O(1)) 算。

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <vector>

using namespace std;

typedef long long ll;
const ll MAXN = 3e5+10;

struct que {
    ll k, id;
    que(){}
    que(ll _k, ll _id){k = _k, id = _id;}
};

struct nod {
    ll l, r, val;
} t[MAXN * 30];

ll N, M, ans[MAXN], siz[MAXN], rt[MAXN], tot;
vector <ll> g[MAXN];
vector <que> q[MAXN];

ll ask(ll, ll, ll, ll, ll);
void dfs(ll, ll);
void dfs2(ll, ll, ll);
void ins(ll&, ll, ll, ll, ll);
ll mer(ll, ll);

int main() {
    scanf("%lld%lld", &N, &M);
    for (ll i = 1, x, y; i < N; i++) {
        scanf("%lld%lld", &x, &y);
        g[x].push_back(y);
        g[y].push_back(x);
    }
    for (ll i = 1, p, k; i <= M; i++) {
        scanf("%lld%lld", &p, &k);
        q[p].push_back(que(k, i));
    }
    dfs(1, 0);
    dfs2(1, 0, 1);
    for (ll i = 1; i <= M; i++)
    	printf("%lld
", ans[i]);
    return 0;
}

void dfs(ll n, ll ff) {
    siz[n] = 1;
    for (unsigned int i = 0 ; i < g[n].size(); i++) {
        ll v = g[n][i];
        if (v == ff) continue;
        dfs(v, n);
        siz[n] += siz[v];
    }
}

void dfs2(ll n, ll ff, ll dep) {
    ins(rt[n], 1, N, dep, siz[n]-1);
    for (unsigned int i = 0; i < g[n].size(); i++) {
        ll v = g[n][i];
        if (v == ff) continue;
        dfs2(v, n, dep+1);
        rt[n] = mer(rt[n], rt[v]);
    }
    for (unsigned int i = 0; i < q[n].size(); i++) {
        ans[q[n][i].id] = min(q[n][i].k, dep-1) * (siz[n]-1) + ask(rt[n], 1, N, dep+1, dep+q[n][i].k);
    }
}

void ins(ll &node, ll l, ll r, ll pos, ll val) {
    if (!node) node = ++tot;
    t[node].val += val;
    if (l == pos && pos == r) return;
    ll mid = (l + r) >> 1;
    if (pos <= mid) ins(t[node].l, l, mid, pos, val);
    else ins(t[node].r, mid+1, r, pos, val);
}

ll ask(ll node, ll l, ll r, ll a, ll b) {
    if (!node) return 0;
    if (a <= l && r <= b) return t[node].val;
    ll mid = (l + r) >> 1, ret = 0;
    if (a <= mid) ret += ask(t[node].l, l, mid, a, b);
    if (mid < b) ret += ask(t[node].r, mid+1, r, a, b);
    return ret;
}

ll mer(ll x, ll y) {
    if (!x || !y) return x | y;
    t[x].val += t[y].val;
    t[x].l = mer(t[x].l, t[y].l);
    t[x].r = mer(t[x].r, t[y].r);
    return x;
}