树链剖分模板

Posted 2021-01-12 zolrk

注意
求子树时，要用子树dfs序连续这个特性，但是不能用如下方法（sub表示x的子树中最大的dfs序）：

void dfs2(int x, int t) {
    top[x] = t;
    dfn[x] = ++cnt;
    subcnt = max(subcnt, cnt);
    rnk[cnt] = x;
    sub[x] = cnt;
    if(!son[x]) return;
    dfs2(son[x], t);
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v != son[x] && v != fa[x]) {
            dfs2(v, v);
        }
    }
    sub[x] = subcnt;
}

因为这样会无法处理叶子节点（直接返回了）
最好的办法是用size，直接算出来
dfn[x] + siz[x] - 1
细节比较多，都在注释里
变量名重复真的很可怕。。。调了半天都没发现

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
using namespace std;
#define debug(x) cerr << #x << "=" << x << endl;
const int MAXN = 100000 + 10;
const int MAXM = 2 * MAXN;

int n,m,ans,vis[MAXN],edge_tot,r,p,dep[MAXN],son[MAXN],top[MAXN],dfn[MAXN],fa[MAXN];
int siz[MAXN],rnk[MAXN],cnt,sub[MAXN],subcnt,last[MAXN],a[MAXN];

struct Edge{
    int u,v,to;
    Edge(){}
    Edge(int u, int v, int to) : u(u), v(v), to(to) {}
}e[MAXM];

inline void add(int u, int v) {
    e[++edge_tot] = Edge(u, v, last[u]);
    last[u] = edge_tot;
}

struct segment{
    int siz,sum,add;
}tr[MAXN * 4];

void down(int now) {
    tr[now*2].sum += tr[now*2].siz * tr[now].add;
    tr[now*2+1].sum += tr[now*2+1].siz * tr[now].add;

    tr[now*2].sum %= p;
    tr[now*2+1].sum %= p;

    tr[now*2].add += tr[now].add;
    tr[now*2+1].add += tr[now].add;

    tr[now].add = 0;
}

void update(int pos) {
    tr[pos].siz = tr[pos*2].siz + tr[pos*2+1].siz;
    tr[pos].sum = tr[pos*2].sum + tr[pos*2+1].sum;
    tr[pos].sum %= p;
}

void build_segment(int now, int l, int r) {
    if(l == r) {
        tr[now].siz = 1;
        tr[now].sum = a[rnk[l]];//求出dfs序对应点的编号，然后才是其值，这里不是a[l]，是l对应的点编号
        return;
    }
    int mid = l+r>>1;
    build_segment(now*2, l, mid);
    build_segment(now*2+1, mid+1, r);
    update(now);
}

void change(int now, int l, int r, int x, int y ,int k) {
    if(x <= l && r <= y) {
        tr[now].add += k;
        tr[now].sum += tr[now].siz * k;
        tr[now].sum %= p;
        return;
    }
    down(now);
    int mid = l+r>>1;
    if(x <= mid)//用不用走左儿子只需要看询问更改边界
        change(now*2, l, mid, x, y, k);
    if(y > mid)
        change(now*2+1, mid+1, r, x, y, k);
    update(now);
}

int query(int now, int l, int r, int x, int y) {
    if(x <= l && r <= y) {
        return tr[now].sum % p;
    }
    down(now);
    long long sum = 0;
    int mid = l+r>>1;
    if(x <= mid)
        sum += query(now*2, l, mid, x, y);
    if(y > mid)
        sum += query(now*2+1, mid+1, r, x, y);
    sum %= p;
    return sum;
}

void dfs1(int x, int f, int depth) {
    fa[x] = f;
    dep[x] = depth;
    siz[x] = 1;
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v == f) continue;//是continue而不是v!=f就执行，因为那样要加大括号。。。
        dfs1(v, x, depth + 1);
        siz[x] += siz[v];
        if(siz[v] > siz[son[x]]) {
            son[x] = v;
        }
    }
}

void dfs2(int x, int t) {
    top[x] = t;
    dfn[x] = ++cnt;
    rnk[cnt] = x;
    if(!son[x]) return;
    dfs2(son[x], t);//保证重链上的点dfs序连续
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v != son[x] && v != fa[x]) {
            dfs2(v, v);
        }
    }
}

int path_query(int x, int y) {
    long long sum = 0;
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);//让x更深，所以先跳x
        sum += query(1, 1, n, dfn[top[x]], dfn[x]);//一定要写topx而不是设tx = topx因为xy会发生交换，提前设置会乱掉。。。
        sum %= p;
        x = fa[top[x]];
    }
    if(dfn[x] > dfn[y]) swap(x, y);
    sum += query(1, 1, n, dfn[x], dfn[y]);
    sum %= p;
    return sum;
}

void path_update(int x, int y, int k) {
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        change(1, 1, n, dfn[top[x]], dfn[x], k);
        x = fa[top[x]];
    }
    if(dfn[x] > dfn[y]) swap(x, y);
    change(1, 1, n, dfn[x], dfn[y], k);
}

int main() {
    scanf("%d%d%d%d", &n, &m, &r, &p);
    for(int i=1; i<=n; i++) {
        scanf("%d", &a[i]);
    }
    for(int i=1; i<n; i++) {
        int x,y;
        scanf("%d%d", &x, &y);
        add(x, y);
        add(y, x);
    }
    dfs1(r, 0, 1);
    dfs2(r,r);
    build_segment(1, 1, n);
    for(int i=1; i<=m; i++) {
        int cmd,x,y,z;
        scanf("%d", &cmd);
        if(cmd == 1) {
            scanf("%d%d%d", &x, &y, &z);
            path_update(x, y, z);
        } else if(cmd == 2) {
            scanf("%d%d", &x, &y);
            printf("%d
", path_query(x, y));
        } else if(cmd == 3) {
            scanf("%d%d", &x, &y);
            change(1, 1, n, dfn[x], dfn[x] + siz[x] - 1, y);
        } else if(cmd == 4) {
            scanf("%d", &x);
            printf("%d
", query(1, 1, n, dfn[x], dfn[x] + siz[x] - 1));
        }
    }
    return 0;
}

两道练手题。。。分别是点修改和边修改
洛谷P3128 [USACO15DEC]最大流Max Flow

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <ctime>
using namespace std;
#define debug(x) cerr << #x << "=" << x << endl;
const int MAXN = 50000 + 10;
const int MAXM = MAXN * 2;

int n,k,dep[MAXN],fa[MAXN],son[MAXN],dfn[MAXN],top[MAXN],rnk[MAXN],dfs_seq,edge_tot,last[MAXN],siz[MAXN];

struct Edge{
    int u,v,to;
    Edge(){}
    Edge(int u, int v, int to) : u(u), v(v), to(to){}
}e[MAXM]; 

inline void add(int u, int v) {
    e[++edge_tot] = Edge(u, v, last[u]);
    last[u] = edge_tot;
}

struct segmentTree{
    int val, tag;
}tr[MAXN*4];

void down(int now) {
    int tagadd = tr[now].tag;
    tr[now*2].val += tagadd;
    tr[now*2+1].val += tagadd;
    tr[now*2].tag += tagadd;
    tr[now*2+1].tag += tagadd;
    tr[now].tag = 0;
}

void change(int now, int l, int r, int cx, int cy, int k) {
    if(cx <= l && r <= cy) {
        tr[now].tag += k;
        tr[now].val += k;
        return;
    }
    down(now);
    int mid = l+r>>1;
    if(cx <= mid) 
        change(now*2, l, mid, cx, cy, k);
    if(cy > mid)
        change(now*2+1, mid+1, r, cx, cy, k);
    tr[now].val = max(tr[now*2].val, tr[now*2+1].val);
}

void dfs1(int x, int father, int depth) {
    dep[x] = depth;
    siz[x] = 1;
    fa[x] = father;
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v == father) continue;
        dfs1(v, x, depth+1);
        siz[x] += siz[v];
        if(siz[v] > siz[son[x]]) son[x] = v;
    }
}

void dfs2(int x, int t) {
    top[x] = t;
    dfn[x] = ++dfs_seq;
    rnk[dfs_seq] = x;
    if(!son[x]) return;
    dfs2(son[x], t);
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v == fa[x] || v == son[x]) continue;
        dfs2(v, v);
    }
}

int path_change(int x, int y, int k) {
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        change(1, 1, n, dfn[top[x]], dfn[x], k);
        x = fa[top[x]];
    }
    
    if(dfn[x] > dfn[y]) swap(x, y);
    change(1, 1, n, dfn[x], dfn[y], k);
}

int main() {
    scanf("%d%d", &n, &k);
    for(int i=1; i<n; i++) {
        int u,v;
        scanf("%d%d", &u, &v);
        add(u, v);
        add(v, u);
    }
    dfs1(1, 0, 1);
    dfs2(1, 1);
    for(int i=1; i<=k; i++) {
        int s, t;
        scanf("%d%d", &s, &t);
        path_change(s, t, 1);
    }
    printf("%d", tr[1].val);
    return 0;
}

洛谷P3038 [USACO11DEC]牧草种植Grass Planting

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <ctime>
using namespace std;
#define debug(x) cerr << #x << "=" << x << endl;
const int MAXN = 100000 + 10;
const int MAXM = MAXN * 2;

int n,m,dep[MAXN],fa[MAXN],son[MAXN],dfn[MAXN],top[MAXN],rnk[MAXN],dfs_seq,edge_tot,last[MAXN],siz[MAXN];

struct Edge{
    int u,v,to;
    Edge(){}
    Edge(int u, int v, int to) : u(u), v(v), to(to){}
}e[MAXM]; 

inline void add(int u, int v) {
    e[++edge_tot] = Edge(u, v, last[u]);
    last[u] = edge_tot;
}

struct segmentTree{
    int val, tag;
}tr[MAXN*4];

void down(int now) {
    tr[now*2].tag += tr[now].tag;
    tr[now*2+1].tag += tr[now].tag;
    tr[now*2].val += tr[now].tag;
    tr[now*2+1].val += tr[now].tag;
    tr[now].tag = 0;
}

void change(int now, int l, int r, int x, int y, int k) {
    if(x <= l && r <= y) {
        tr[now].tag += k;
        tr[now].val += k;
        return;
    }
    down(now);
    int mid = l+r>>1;
    if(x <= mid) 
        change(now*2, l, mid, x, y, k);
    if(y > mid)
        change(now*2+1, mid+1, r, x, y, k);
}

int query(int now, int l, int r, int x) {
    if(l == r) return tr[now].val;
    down(now);
    int mid = l+r>>1;
    if(x <= mid)
        return query(now*2, l, mid, x);
    else
        return query(now*2+1, mid+1, r, x);
}

void dfs1(int x, int father, int depth) {
    dep[x] = depth;
    siz[x] = 1;
    fa[x] = father;
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v == father) continue;
        dfs1(v, x, depth+1);
        siz[x] += siz[v];
        if(siz[v] > siz[son[x]]) son[x] = v;
    }
}

void dfs2(int x, int t) {
    top[x] = t;
    dfn[x] = ++dfs_seq;
    rnk[dfs_seq] = x;
    if(!son[x]) return;
    dfs2(son[x], t);
    for(int i=last[x]; i; i=e[i].to) {
        int v = e[i].v;
        if(v == fa[x] || v == son[x]) continue;
        dfs2(v, v);
    }
}

void path_change(int x, int y, int k) {
    while(top[x] != top[y]) {
        if(dep[top[x]] < dep[top[y]]) swap(x, y);
        change(1, 1, n, dfn[top[x]], dfn[x], k);
        x = fa[top[x]];
    }
    if(dfn[x] > dfn[y]) swap(x, y);
    change(1, 1, n, dfn[x], dfn[y], k);
    change(1, 1, n, dfn[x], dfn[x], -k);
}

int main() {
    scanf("%d%d", &n, &m);
    for(int i=1; i<n; i++) {
        int u,v;
        scanf("%d%d", &u, &v);
        add(u, v);
        add(v, u);
    }
    dfs1(1, 0, 1);
    dfs2(1, 1);
    for(int i=1; i<=m; i++) {
        char cmd;
        cin >> cmd;
        int ai, bi;
        scanf("%d%d", &ai, &bi);
        if(cmd == 'P') {
            path_change(ai, bi, 1);
        } else {
            printf("%d
", query(1, 1, n, max(dfn[ai], dfn[bi])));
        }
    }
    return 0;
}