Query on the Subtree

题目描述

Bobo has a tree, whose vertices are conveniently labeled by $1, 2, \ldots, n$. At the very begining, the $i$-th vertex is assigned with weight $w_i$.

There are $q$ operations. Each operations are of the following $2$ types:

• change the weight of vertex $v$ into $x$ (denoted as “! $v$ $x$”),
• ask the total weight of vertices whose distance are no more than $d$ away from vertex $v$ (denoted as “? $v$ $d$”).

Note that the distance between vertex $u$ and $v$ is the number of edges on the shortest path between them.

题意概述

给定一棵有$n$个节点的树，第$i$个节点的权值为$w_i$。有两种操作：① 将节点$v$的权值变成$x$；② 询问与节点$v$的距离不超过$d$的所有节点的权值和。共有$q$次操作。

数据范围：$1 \le n, q \le 10^5, \ 0 \le w_i \le 10^4, \ 0 \le x \le 10^4, \ 0 \le d \le n$。

算法分析

对于我们枚举的某一个重心$x$，从它子树的重心向它连一条“虚”边，这样就形成了一棵由重心相连构成的“重心树”。由重心的性质可得，这棵树最多只有$O(\log n)$层。

如图，黑色表示原树上的边，红色表示重心树上的“虚”边。

对于每个节点，我们用树状数组存下它在“重心树”上的子树节点到它的距离为$p$（在原树上的距离）的权值和。当我们询问到节点$v$的距离不超过$d$的节点的权值和时，答案就等于$v$的树状数组中$d$的前缀和，再加上“重心树”上$v$的祖先节点$u$的树状数组中$d-dist_{u, v}$的前缀和。可以发现，$u$在“重心树”上包含$v$的子树中的节点被重复计算了。

如图，在计算到节点$3$的距离不超过$3$的节点的权值和时，计算了节点$3$子树中的节点$1, 3, 6, 8$；在计算到节点$3$的祖先节点$2$的距离不超过$3-2=1$的节点的权值和时，节点$1$又被计算了一次。

根据容斥原理，只需对每个节点再用一个树状数组记录下它在“重心树”上的子树节点到它在“重心树”上的父节点的距离为$p$（在原树上的距离）的权值和，每次计算时减去这个树状数组中$d-dist_{u, v}$的前缀和，就得到了正确答案。

修改节点$v$的权值时，只需更新“重心树”上$v$所有祖先节点的树状数组即可。

树上两点间的距离可以通过倍增求 LCA 得到，不过有更简便的方法。由于每个节点只会被$O(\log n)$个重心搜到，因此可以存下每个节点到其所有祖先重心节点的距离。

在“重心树”上操作的时间复杂度为$O(\log n)$，树状数组的时间复杂度为$O(\log n)$，总时间复杂度为$O(q\log^2n)$。

“重心树”上每一层节点的树状数组空间复杂度为$O(n)$，有$O(\log n)$层，每个节点上要存$O(\log n)$个距离，总空间复杂度为$O(n\log n)$。

代码

#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
struct edge {
    int v, nxt;
} e[200001];
struct binary_indexed_tree {
    int n;
    vector<int> a;
    void init(int size) { n = size + 10, a.clear(), a.resize(n); }
    void add(int p, int t) {
        if (p) for (int i = p; i < n; i += i & -i) a[i] += t;
    }
    int sum(int p) {
        if (p <= 0) return 0;
        int ret = 0;
        for (int i = min(p, n - 1); i; i -= i & -i) ret += a[i];
        return ret;
    }
} tree[200001];
int n, q, nume, tot, top, root, h[100001], w[100001];
int size[100001], f[100001], up[100001], id[100001][2];
bool vis[100001];
vector<int> dist[100001];
void add_edge(int u, int v) {
    e[++nume].v = v, e[nume].nxt = h[u], h[u] = nume;
    e[++nume].v = u, e[nume].nxt = h[v], h[v] = nume;
}
void get_root(int t, int fa) {
    size[t] = 1, f[t] = 0;
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v] && e[i].v != fa) {
            get_root(e[i].v, t);
            size[t] += size[e[i].v], f[t] = max(f[t], size[e[i].v]);
        }
    }
    f[t] = max(f[t], tot - size[t]);
    if (f[t] < f[root]) root = t;
}
void get_dist(int t, int fa, int depth, int flag) {
    if (flag) dist[t].push_back(depth);
    else {
        tree[id[root][0]].add(depth, w[t]);
        if (dist[t].size() > 1) tree[id[root][1]].add(dist[t][dist[t].size() - 2], w[t]);
    }
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v] && e[i].v != fa) get_dist(e[i].v, t, depth + 1, flag);
    }
}
void solve(int t) {
    vis[t] = true;
    for (int i = h[t]; i; i = e[i].nxt) if (!vis[e[i].v]) get_dist(e[i].v, t, 1, 0);
    for (int i = h[t]; i; i = e[i].nxt) {
        if (!vis[e[i].v]) {
            root = 0, tot = size[e[i].v], get_root(e[i].v, t), up[root] = t;
            tree[id[root][0] = ++top].init(size[e[i].v]);
            tree[id[root][1] = ++top].init(size[e[i].v]);
            tree[id[root][1]].add(dist[root][dist[root].size() - 1], w[root]);
            get_dist(root, 0, 0, 1);
            solve(root);
        }
    }
}
void modify(int t, int d) {
    int p = t, top = dist[t].size() - 1;
    while (p) {
        if (top >= 0) tree[id[p][0]].add(dist[t][top], d);
        if (top > 0) tree[id[p][1]].add(dist[t][top - 1], d);
        p = up[p], --top;
    }
}
int ask(int t, int d) {
    int p = t, ret = 0, top = dist[t].size() - 1;
    while (p) {
        int len;
        if (top >= 0) len = d - dist[t][top];
        if (top >= 0 && len >= 0) ret += tree[id[p][0]].sum(len) + w[p];
        if (top > 0) len = d - dist[t][top - 1];
        if (top > 0 && len >= 0 && up[p]) ret -= tree[id[p][1]].sum(len);
        p = up[p], --top;
    }
    return ret;
}
int main()
{
    while (scanf("%d%d", &n, &q) != EOF) {
        nume = top = 0;
        memset(up, 0, sizeof(up));
        memset(h, 0, sizeof(h));
        memset(vis, 0, sizeof(vis));
        for (int i = 1; i <= n; ++i) {
            scanf("%d", &w[i]);
            dist[i].clear();
        }
        for (int i = 1; i < n; ++i) {
            int u, v;
            scanf("%d%d", &u, &v), add_edge(u, v);
        }
        root = 0, tot = f[0] = n, get_root(1, 0);
        tree[id[root][0] = ++top].init(size[1]);
        tree[id[root][1] = ++top].init(size[1]);
        get_dist(root, 0, 0, 1);
        solve(root);
        while (q--) {
            char oper = ' ';
            while (oper != '!' && oper != '?') scanf("%c", &oper);
            int v, d;
            scanf("%d%d", &v, &d);
            if (oper == '!') modify(v, d - w[v]), w[v] = d;
            else printf("%d\n", ask(v, d));
        }
    }
    return 0;
}