牛客 树的距离

Posted 2021-09-11 TURNINING

传送门

题意：有一颗有根树，根节点为1，赋有边权。现有m次询问，每次询问在x的子树中所有与x距离大于等于k的点与x的距离之和。

思路1：x与子树中某节点的距离就是 该节点到根节点的距离 - x到根节点的距离。那么我们维护一下x子树中所有点到根节点的距离，求出所有距离大于等于$k+d[x]$的和$sum$，再求出点的个数$cnt$，答案就是$sum-d[x]\ast cnt。$显然主席树可行。这是一种在线的做法

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

typedef long long ll;

#define endl "\\n"
#define lsn(u) (tr[u].ls)
#define rsn(u) (tr[u].rs)
#define mid (l + r >> 1)

const int MAXN = 2e5 + 10;

struct Node {
    int ls, rs;
    ll sum, cnt;
}tr[10000000];

struct edge {
    int to, cost;
};

vector<edge> g[MAXN];

int n, m;
int tot, num, rt[MAXN];
int in[MAXN], out[MAXN];
ll d[MAXN], cnt, sum;

int update(int u, ll l, ll r, ll pos) {
    int cur = ++tot;
    tr[cur] = tr[u];
    tr[cur].cnt++;
    tr[cur].sum += pos;
    if(l == r) return cur;
    if(pos <= mid) 
        lsn(cur) = update(lsn(cur), l, mid, pos);
    else 
        rsn(cur) = update(rsn(cur), mid+1, r, pos);
    return cur;
}   

void query(int u, int v, ll l, ll r, ll L, ll R) {
    if(L <= l && r <= R) {
        cnt += tr[v].cnt - tr[u].cnt;
        sum += tr[v].sum - tr[u].sum;
    } 
    else {
        if(R <= mid)  query(lsn(u), lsn(v), l, mid, L, R);
        else if(L > mid)  query(rsn(u), rsn(v), mid+1, r, L, R);
        else query(lsn(u), lsn(v), l, mid, L, R), query(rsn(u), rsn(v), mid+1, r, L, R);
    }
}

void dfs(int u, int p) {
    in[u] = ++num;
    rt[in[u]] = update(rt[in[u]-1], 1, (ll)1e12, d[u]);
    for(auto e : g[u]) {
        if(e.to == p) continue;
        d[e.to] = d[u] + e.cost;
        dfs(e.to, u);
    }
    out[u] = num;
}

void solve() {
    cin >> n;
    for(int i = 1; i < n; i++) {
        int u, c; cin >> u >> c;
        g[u].emplace_back(edge{i+1, c});
        g[i+1].emplace_back(edge{u, c});
    }
    d[1] = 1;
    dfs(1, 0);
    cin >> m;
    while(m--) {
        cnt = sum = 0;
        ll x, k; cin >> x >> k;
        query(rt[in[x]-1], rt[out[x]], 1, (ll)1e12, d[x]+k, (ll)1e12);
        if(cnt == 0) cout << 0 << endl;
        else cout << sum - cnt * d[x] << endl;
    }
}

int main() {
    ios::sync_with_stdio(false);
    solve();
    return 0;
}

思路2：考虑离线怎么做。对于每个查询，我们其实只要知道大于等于d[x]+k的数量和他们的和。我们把询问按d[x]+k从大到小排序，然后再把每个节点按d[x]排序。对于每次的查询，我们只将大于等于d[x]+k的节点加入，这样每个节点就只会被加入一次。求和 和 求数量可以用树状数组简单维护一下。
这种方法更值得学习。

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

typedef long long ll;
typedef pair<ll, ll> P;

#define endl "\\n"

const int MAXN = 2e5 + 10;

struct edge {
    int to, cost;
};

struct Q {         
    int l, r, id;
    ll dis, w;
    bool operator < (const Q &b) const {
        return dis > b.dis;
    }
}q[MAXN];

struct Node {
    int pos;
    ll w;
    bool operator < (const Node &b) const {
        return w > b.w;
    }
}a[MAXN];

vector<edge> g[MAXN];

ll cnt[MAXN], sum[MAXN], d[MAXN], ans[MAXN];
int n, m, tot;
int in[MAXN], out[MAXN];

void dfs(int u, int p) {
    in[u] = ++tot;
    for(auto e : g[u]) {
        if(e.to == p) continue;
        d[e.to] = d[u] + e.cost;
        dfs(e.to, u);
    }
    out[u] = tot;
}

void upd(int x, ll w) {
    while(x <= n) {
        cnt[x] += 1;
        sum[x] += w;
        x += x & -x;
    }
}

P ask(int x) {
    ll res1 = 0, res2 = 0;
    while(x > 0) {
        res1 += cnt[x];
        res2 += sum[x];
        x -= x & -x;
    }  
    return P(res1, res2);
}

void solve() {
    cin >> n;
    for(int i = 1; i < n; i++) {
        int v, w; cin >> v >> w;
        g[v].emplace_back(edge{i+1, w});
        g[i+1].emplace_back(edge{v, w});
    }
    dfs(1, 0);
    cin >> m;
    for(int i = 1; i <= m; i++) {
        int x; ll w; cin >> x >> w;
        q[i].l = in[x]; q[i].r = out[x];
        q[i].w = d[x]; q[i].dis = w + d[x];
        q[i].id = i;
    }
    for(int i = 1; i <= n; i++) {
        a[i].pos = in[i];
        a[i].w = d[i];
    }
    sort(q+1, q+1+m);
    sort(a+1, a+1+n);
    for(int i = 1, j = 1; i <= m; i++) {
        while(j <= n && a[j].w >= q[i].dis) {
            upd(a[j].pos, a[j].w);
            j++;
        }
        P p1 = ask(q[i].r), p2 = ask(q[i].l - 1);
        ll cnt = p1.first - p2.first, sum = p1.second - p2.second;

        if(cnt == 0) ans[q[i].id] = 0;
        else ans[q[i].id] = sum - cnt * q[i].w;
    }
    for(int i = 1; i <= m; i++) {
        cout << ans[i] << endl; 
    }
}

int main() {
    ios::sync_with_stdio(false);
    solve();
    return 0;
}