poj 2763(LCA + dfs序 +树状数组)

Posted 2021-01-18 sugewud

算是模板题了

可以用dfs序维护点到根的距离

注意些LCA的时候遇到MAXM，要-1

#include<cstdio>
#include<algorithm>
#include<cstring>
#define REP(i, a, b) for(register int i = (a); i < (b); i++)
#define _for(i, a, b) for(register int i = (a); i <= (b); i++)
using namespace std;

const int MAXN = 1e5 + 10;
const int MAXM = 20;
struct Edge{ int to, next, w; } e[MAXN << 1];
int head[MAXN], tot;

int U[MAXN], V[MAXN], W[MAXN], d[MAXN];
int L[MAXN], R[MAXN], f[MAXN], cnt;
int up[MAXN][MAXM+10], n, s, q;

void AddEdge(int from, int to, int w)
{
    e[tot] = Edge{to, head[from], w};
    head[from] = tot++;
}

inline int lowbit(int x) { return x & (-x); }

void add(int x, int p)
{
    for(; x <= n; x += lowbit(x))
        f[x] += p;
}

int sum(int x)
{
    int res = 0;
    for(; x; x -= lowbit(x))
        res += f[x];
    return res;
}

inline void modify(int u, int w)
{
    add(L[u], w); add(R[u] + 1, -w);
}

void dfs(int u, int fa, int op)
{
    if(!op) L[u] = ++cnt;
    for(int i = head[u]; ~i; i = e[i].next)
    {
        int v = e[i].to;
        if(v == fa) continue;
        if(!op)
        {        
            up[v][0] = u;
            d[v] = d[u] + 1;
        }
        if(op) modify(v, e[i].w);
        dfs(v, u, op);
    }
    if(!op) R[u] = cnt;
}

void init()
{
    up[1][0] = 1;
    REP(j, 1, MAXM)
        _for(i, 1, n)
            up[i][j] = up[up[i][j-1]][j-1];
}

int lca(int u, int v)
{
    if(d[u] < d[v]) swap(u, v);
    for(int j = MAXM - 1; j >= 0; j--)
        if(d[up[u][j]] >= d[v])
            u = up[u][j];
    if(u == v) return u;
    for(int j = MAXM - 1; j >= 0; j--)
        if(up[u][j] != up[v][j])
            u = up[u][j], v = up[v][j];
    return up[u][0];
}

inline int dis(int u, int v)
{
    return sum(L[u]) + sum(L[v]) - 2 * sum(L[lca(u, v)]);
}

int main()
{
    while(~scanf("%d%d%d", &n, &q, &s))
    {
        memset(head, -1, sizeof(head)); 
        memset(f, 0, sizeof(f));
        tot = cnt = 0;
        
        REP(i, 1, n)
        {
            int u, v, w;
            scanf("%d%d%d", &u, &v, &w); 
            U[i] = u; V[i] = v; W[i] = w;
            AddEdge(u, v, w); AddEdge(v, u, w);
        }
    
        dfs(1, -1, 0);
        dfs(1, -1, 1);
        init();
    
        while(q--)
        {
            int op, x, y;
            scanf("%d", &op);
            if(op == 0) 
            {
                scanf("%d", &x);
                printf("%d
", dis(x, s));
                s = x;
            }
            else 
            {
                scanf("%d%d", &x, &y);
                int u = U[x], v = V[x], w = W[x];
                if(up[u][0] == v) swap(u, v);
                modify(v, y - w);
                W[x] = y;
            }
        }
    }

    return 0;
}