Luogu 2590 [ZJOI2008]樹的統計 / HYSBZ 1036 [ZJOI2008]樹的統計Count （樹鏈剖分，LCA，線段樹）

Luogu 2590 [ZJOI2008]樹的統計 / HYSBZ 1036 [ZJOI2008]樹的統計Count （樹鏈剖分，LCA，線段樹）

Description

一棵樹上有n個節點，編號分別爲1到n，每個節點都有一個權值w。我們將以下面的形式來要求你對這棵樹完成一些操作：
I. CHANGE u t : 把結點u的權值改爲t

II. QMAX u v: 詢問從點u到點v的路徑上的節點的最大權值

III. QSUM u v: 詢問從點u到點v的路徑上的節點的權值和 注意：從點u到點v的路徑上的節點包括u和v本身

Input

輸入文件的第一行爲一個整數n，表示節點的個數。
接下來n – 1行，每行2個整數a和b，表示節點a和節點b之間有一條邊相連。
接下來一行n個整數，第i個整數wi表示節點i的權值。
接下來1行，爲一個整數q，表示操作的總數。
接下來q行，每行一個操作，以“CHANGE u t”或者“QMAX u v”或者“QSUM u v”的形式給出。

Output

對於每個“QMAX”或者“QSUM”的操作，每行輸出一個整數表示要求輸出的結果。

Sample Input

4
1 2
2 3
4 1
4 2 1 3
12
QMAX 3 4
QMAX 3 3
QMAX 3 2
QMAX 2 3
QSUM 3 4
QSUM 2 1
CHANGE 1 5
QMAX 3 4
CHANGE 3 6
QMAX 3 4
QMAX 2 4
QSUM 3 4

Sample Output

4
1
2
2
10
6
5
6
5
16

Http

Luogu:https://www.luogu.org/problem/show?pid=2590
HYSBZ:http://www.lydsy.com/JudgeOnline/problem.php?id=1036

Source

樹鏈剖分，最近公共祖先LCA，線段樹

解決思路

這是一道樹鏈剖分入門題。
考慮將樹按照子樹大小剖分成輕鏈和重鏈，然後用線段樹來維護區間和和區間最大值。
需要注意的是，在樹上進行鏈的跳轉時，使用原編號，而若要用線段樹查詢或修改線段樹時，要使用線段樹中的編號（即第2次dfs中給點新分配的編號）

代碼

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
using namespace std;

#define ll long long

const int maxN=350000;
const int inf=2147483647;

class Edge//邊
{
public:
    int u,v,nex;
};

class SegmentTree//線段樹
{
public:
    ll mx,sum;//維護最大值和區間和
    SegmentTree()
        {
            mx=0;
            sum=0;
        }
};

class Tree_Chain//樹鏈剖分中需要維護的一些值
{
public:
    int f,sz,top,hson,depth;//父節點，大小，向上跳的最高位置，重兒子，深度
};

int n;
//Tree 原樹的一些數據
int cnt;
int Head[maxN];
Edge E[maxN*2];
ll W[maxN];//點權
Tree_Chain T[maxN];//樹鏈剖分維護的值
int dfn[maxN];//記錄樹鏈剖分後的點的新編號（按照先重鏈再輕鏈）
SegmentTree Seg[maxN*4];//線段樹

void Add_Edge(int u,int v);
//Tree_Chain
void dfs1(int u,int father);//第一遍dfs求出T中的一些數據
void dfs2(int u,int Top);//第二遍dfs構造Top和分配新編號
ll LCA_sum(int u,int v);//求解區間和
ll LCA_max(int u,int v);//求解區間最大值
//SegmentTree
void Update(int l,int r,int now,int pos,int key);//線段樹單點更新
ll Query_sum(int l,int r,int now,int ql,int qr);//查詢區間和
ll Query_max(int l,int r,int now,int ql,int qr);//查詢區間最大值

int main()
{
    cnt=0;
    memset(Head,-1,sizeof(Head));
    scanf("%d",&n);
    for (int i=1;i<n;i++)//輸入
    {
        int u,v;
        scanf("%d%d",&u,&v);
        Add_Edge(u,v);
        Add_Edge(v,u);
    }
    T[1].f=0;//默認1爲根節點
    T[1].depth=1;
    dfs1(1,1);//第一遍dfs
    cnt=0;//標號置爲0
    dfs2(1,1);//第二遍dfs
    for (int i=1;i<=n;i++)//輸入點權
        scanf("%lld",&W[i]);
    for (int i=1;i<=n;i++)
        Update(1,n,1,dfn[i],W[i]);//將點權按照新編號放入線段樹
    int qus;
    scanf("%d",&qus);
    while (qus--)//處理操作
    {
        char str[10];
        scanf("%s",str);
        if (str[0]=='Q')
        {
            int u,v;
            scanf("%d%d",&u,&v);
            if (str[1]=='M')
                printf("%lld\n",LCA_max(u,v));//查詢最大值
            else
                printf("%lld\n",LCA_sum(u,v));//查詢和
        }
        else
        {
            int pos,key;
            scanf("%d%d",&pos,&key);
            pos=dfn[pos];
            Update(1,n,1,pos,key);//單點修改
        }
    }
    return 0;
}

void Add_Edge(int u,int v)//在圖中加一條邊
{
    cnt++;
    E[cnt].nex=Head[u];
    Head[u]=cnt;
    E[cnt].u=u;
    E[cnt].v=v;
    return;
}

void dfs1(int u,int father)//第一遍dfs
{
    T[u].sz=1;//初始化基本信息
    T[u].hson=0;
    for (int i=Head[u];i!=-1;i=E[i].nex)
    {
        int v=E[i].v;
        if (v==father)
            continue;
        T[v].f=u;//將子節點信息更新
        T[v].depth=T[u].depth+1;
        dfs1(v,u);
        T[u].sz=T[u].sz+T[v].sz;
        if ((T[u].hson==0)||(T[T[u].hson].sz<T[v].sz))//求出重兒子
            T[u].hson=v;
    }
    return;
}

void dfs2(int u,int Top)//第二遍dfs，Top是當前傳下來的能到達的重鏈頂端，當u==Top時，爲輕鏈
{
    cnt++;
    dfn[u]=cnt;//分配新編號
    T[u].top=Top;//分配Top
    if (T[u].hson==0)//沒有重兒子，說明到了葉子節點
        return;
    dfs2(T[u].hson,Top);//讓重兒子繼承重鏈
    for (int i=Head[u];i!=-1;i=E[i].nex)
    {
        int v=E[i].v;
        if ((v==T[u].f)||(v==T[u].hson))
            continue;
        dfs2(v,v);//其他兒子開輕鏈
    }
    return;
}

ll LCA_max(int u,int v)//求解u到v路徑上的最大點權
{
    ll mx=-inf;
    while (T[u].top!=T[v].top)
    {
        if (T[T[u].top].depth<T[T[v].top].depth)//保證每一次u的深度儘量大
            swap(u,v);
        mx=max(mx,Query_max(1,n,1,dfn[T[u].top],dfn[u]));//讓深度大的向上跳，注意這裏要轉成線段樹中的編號
        u=T[T[u].top].f;
    }
    if (dfn[v]<dfn[u])//保證u的深度更小
        swap(u,v);
    mx=max(mx,Query_max(1,n,1,dfn[u],dfn[v]));//查詢鏈
    return mx;
}

ll LCA_sum(int u,int v)//求解u到v上的點權之和，與上面基本一樣
{
    ll sum=0;
    while (T[u].top!=T[v].top)
    {
        if (T[T[u].top].depth<T[T[v].top].depth)
            swap(u,v);
        sum=sum+Query_sum(1,n,1,dfn[T[u].top],dfn[u]);
        u=T[T[u].top].f;
    }
    if (dfn[v]<dfn[u])
        swap(u,v);
    sum=sum+Query_sum(1,n,1,dfn[u],dfn[v]);
    return sum;
}

void Update(int l,int r,int now,int pos,int key)//線段樹單點修改
{
    if (l==r)
    {
        Seg[now].mx=Seg[now].sum=key;
        return;
    }
    int mid=(l+r)/2;
    if (pos<=mid)
        Update(l,mid,now*2,pos,key);
    if (pos>=mid+1)
        Update(mid+1,r,now*2+1,pos,key);
    Seg[now].sum=Seg[now*2].sum+Seg[now*2+1].sum;
    Seg[now].mx=max(Seg[now*2].mx,Seg[now*2+1].mx);
    return;
}

ll Query_sum(int l,int r,int now,int ql,int qr)//線段樹查詢區間和
{
    if ((l==ql)&&(r==qr))
        return Seg[now].sum;
    int mid=(l+r)/2;
    if (qr<=mid)
        return Query_sum(l,mid,now*2,ql,qr);
    else if (ql>=mid+1)
        return Query_sum(mid+1,r,now*2+1,ql,qr);
    else
        return Query_sum(l,mid,now*2,ql,mid)+Query_sum(mid+1,r,now*2+1,mid+1,qr);
}

ll Query_max(int l,int r,int now,int ql,int qr)//線段樹查詢區間最大值
{
    if ((l==ql)&&(r==qr))
        return Seg[now].mx;
    int mid=(l+r)/2;
    if (qr<=mid)
        return Query_max(l,mid,now*2,ql,qr);
    if (ql>=mid+1)
        return Query_max(mid+1,r,now*2+1,ql,qr);
    return max(Query_max(l,mid,now*2,ql,mid),Query_max(mid+1,r,now*2+1,mid+1,qr));
}