hdu 6162 Ch’s gift

Problem

acm.hdu.edu.cn/showproblem.php?pid=6162

Meaning

一棵 n 个点的树，每个节点有个 $c_{i}$。m 次询问，每次求在点 s 到点 t 的最短路之间，$a\leq c_{i}\leq b$ 的 $c_{i}$ 的总和。

Analysis

任意找一个点做树根，对每条树上的路径都建一棵主席树（结点记录从根到当前结点这条路径的 $c_{i}$ 的序列里，值域在 [ L , R ] 的 $c_{i}$ 的和）。
对每一个询问，找出 s 和 t 的 LCA，然后query(s) + query(t) - 2 * query(LCA) + x，其中的 x 当 $a\leq c[LCA]\leq b$ 时为 c[LCA]，否则 x 为 0。

Code

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 1e5, LOG_N = 20, C = 1e9, LOG_C = 31;

int head[N+1], to[N<<1], nxt[N<<1];

void add_edge(int f, int t, int id)
{
    to[id] = t;
    nxt[id] = head[f];
    head[f] = id;
}

/*--- LCA ---*/
int depth[N+1], pa[N+1][LOG_N];

void dfs(int v, int f, int d)
{
    depth[v] = d;
    pa[v][0] = f;
    for(int i = head[v]; ~i; i = nxt[i])
        if(to[i] != f)
            dfs(to[i], v, d + 1);
}

void pre_lca(int n, int rt)
{
    dfs(rt, -1, 0);
    for(int k = 0; 1 << k + 1 < n; ++k)
        for(int v = 1; v <= n; ++v)
            if(pa[v][k] < 0)
                pa[v][k+1] = -1;
            else
                pa[v][k+1] = pa[pa[v][k]][k];
}

int lca(int x, int y)
{
    if(depth[x] > depth[y])
        swap(x, y);
    for(int i = 0, j = depth[y] - depth[x]; j > 0; j >>= 1, ++i)
        if(j & 1)
            y = pa[y][i];
    if(x == y)
        return x;
    for(int k = LOG_N - 1; k >= 0; --k)
        if(pa[x][k] != pa[y][k])
        {
            x = pa[x][k];
            y = pa[y][k];
        }
    return pa[x][0];
}

/*--- 主席树 ---*/
struct node
{
    int l, r;
    long long sum;
} tree[N * LOG_C];
int root[N+1], sz;

int add_link(int v, int l, int r, int pre)
{
    int rt = ++sz;
    tree[rt] = tree[pre];
    tree[rt].sum += v;
    if(l != r)
    {
        int m = l + r >> 1;
        if(m < v)
            tree[rt].r = add_link(v, m+1, r, tree[pre].r);
        else
            tree[rt].l = add_link(v, l, m, tree[pre].l);
    }
    return rt;
}

int c[N+1], big;

void build(int v, int f, int pre)
{
    root[v] = add_link(c[v], 1, big, pre);
    for(int i = head[v]; ~i; i = nxt[i])
        if(to[i] != f)
            build(to[i], v, root[v]);
}

long long query(int ql, int qr, int l, int r, int id)
{
    if(ql <= l && r <= qr)
        return tree[id].sum;
    if(r < ql || qr < l)
        return 0;
    int m = l + r >> 1;
    return query(ql, qr, l, m, tree[id].l) +
        query(ql, qr, m+1, r, tree[id].r);
}

int main()
{
    tree[0].sum = tree[0].l = tree[0].r = 0;
    int n, m;
    while(~scanf("%d%d", &n, &m))
    {
        big = 0; // 最大的 ci
        for(int i = 1; i <= n; ++i)
        {
            scanf("%d", c+i);
            if(c[i] > big)
                big = c[i];
        }
        // 建图
        memset(head, -1, sizeof head);
        for(int i = 1, f, t, id = 0; i < n; ++i)
        {
            scanf("%d%d", &f, &t);
            add_edge(f, t, id++);
            add_edge(t, f, id++);
        }
        // 深搜建树
        sz = 0;
        build(1, -1, 0);
        // LCA 预处理
        pre_lca(n, 1);
        // 询问
        long long ans, x, y, z;
        for(int s, t, a, b, anc; m--; )
        {
            scanf("%d%d%d%d", &s, &t, &a, &b);
            anc = lca(s, t);

            x = query(a, b, 1, big, root[s]);
            y = query(a, b, 1, big, root[t]);
            z = query(a, b, 1, big, root[anc]);

            ans = x + y - z - z;
            if(a <= c[anc] && c[anc] <= b)
                ans += c[anc];

            printf("%I64d%c", ans, m ? ' ' : '\n');
        }
    }
    return 0;
}