Aladdin and the Return Journey

#include <bits/stdc++.h>
using namespace std;
#define Size 100005

#define root 1
#define N 30100
#define LN 14

vector<int> Graph[N];
int treeArray[N], Node;
int chainNo, chainInd[N], chainHead[N], posInTree[N];
int depth[N], sparseTable[N][LN], subSize[N];
long long tree[N * 6];
long long Cost[N];
int n;

void build_tree(int cur, int Start, int End) {
    if (Start == End) {
        tree[cur] = treeArray[Start];
        return;
    }
    int left = cur*2;
    int right = left+1;
    int mid = (Start+End)/2;
    build_tree(left,Start,mid);
    build_tree(right,mid+1,End);
    tree[cur] = tree[left]+tree[right];
}

void update_tree(int cur, int Start, int End, int u, int v, long long value) {
    if (Start > v || End < u)
        return;
    if (Start >= u && End <= v) {
        tree[cur] = value;
        return;
    }
    int left = cur*2;
    int right = left+1;
    int mid = (Start+End)/2;

    update_tree(left,Start,mid,u,v,value);
    update_tree(right,mid+1,End,u,v,value);
    tree[cur] = tree[left]+tree[right];
}

long long query_Tree(int cur, int Start, int End, int u, int v) {
    if (Start > v || End < u) {
        return 0;
    }
    if (Start >= u && End <= v) {
        return tree[cur];
    }
    int left = cur*2;
    int right = left+1;
    int mid = (Start+End)/2;

    long long v1 = query_Tree(left,Start,mid,u,v);
    long long v2 = query_Tree(right,mid+1,End,u,v);
    long long sum = v1+v2;
    //printf("Cur sum: %lld\n",sum);
    return sum;
}

long long query_up(int u, int v) {
    if (u == v) {
        //return 0;
        return query_Tree(1, 1, Node, posInTree[u], posInTree[u]);
    }
    int uchain, vchain = chainInd[v];
    long long ans = 0;
    while (true) {
        uchain = chainInd[u];
        if (uchain == vchain) {
            int st = posInTree[v];
            int ed = posInTree[u];
            long long ret = query_Tree(1, 1, Node, st, ed);
            //printf("st: %d , ed: %d = %lld\n",st,ed,ret);
            ans += ret;
            break;
        }
        int st = posInTree[chainHead[uchain]];
        int ed = posInTree[u];
        long long ret = query_Tree(1, 1, Node, st, ed);
        //printf("st: %d , ed: %d = %lld\n",st,ed,ret);
        ans += ret;
        u = chainHead[uchain];
        u = sparseTable[u][0];
        //cout << ans << endl;
    }
    return ans;
}

int LCA(int u, int v) {
    if (depth[u] < depth[v]){
        swap(u, v);
    }
    int diff = depth[u] - depth[v];
    for (int i = 0; i < LN; i++) {
        if ((diff >> i) & 1) {
            u = sparseTable[u][i];
        }
    }
    if (u == v){
        return u;
    }
    for (int i = LN - 1; i >= 0; i--) {
        if (sparseTable[u][i] != sparseTable[v][i]) {
            u = sparseTable[u][i];
            v = sparseTable[v][i];
        }
    }
    return sparseTable[u][0];
}

void build_Sparse_Table() {
    for (int i = 1; i < LN; i++) {
        for (int j = 1; j <= N; j++) {
            if (sparseTable[j][i - 1] != -1) {
                sparseTable[j][i] = sparseTable[sparseTable[j][i - 1]][i - 1];
            }
        }
    }
}

long long query(int u, int v) {
    //return query_up(11,1);
    int lca = LCA(u, v);
    long long ret1 = query_up(u, lca);
    long long ret2 = query_up(v, lca);
    long long merge = query_up(lca, lca);
    //printf("LCA: %d , ret1: %lld , ret2: %lld\n",lca,ret1,ret2);
    return (ret1+ret2-merge);
}

void updateNode(int u, long long val) {
    update_tree(1, 1, Node, posInTree[u],posInTree[u], val);
}

void HLD(int curNode,int prev) {
    if (chainHead[chainNo] == -1) {
        chainHead[chainNo] = curNode;
    }
    chainInd[curNode] = chainNo;
    posInTree[curNode] = Node;
    treeArray[Node++] = Cost[curNode];

    int newNode = -1;
    for (int i = 0; i < (int) Graph[curNode].size(); i++) {
        if (Graph[curNode][i] != prev) {
            if (newNode == -1 || subSize[newNode] < subSize[Graph[curNode][i]]) {
                newNode = Graph[curNode][i];
            }
        }
    }

    if (newNode != -1) {
        HLD(newNode, curNode);
    }

    for (int i = 0; i < (int) Graph[curNode].size(); i++) {
        if (Graph[curNode][i] != prev) {
            if (newNode != Graph[curNode][i]) {
                chainNo++;
                HLD(Graph[curNode][i], curNode);
            }
        }
    }
}

void DFS(int cur, int prev, int level = 0) {
    sparseTable[cur][0] = prev;
    depth[cur] = level;
    subSize[cur] = 1;
    for (int i = 0; i < (int) Graph[cur].size(); i++)
        if (Graph[cur][i] != prev) {
            DFS(Graph[cur][i], cur, level + 1);
            subSize[cur] += subSize[Graph[cur][i]];
        }
}

void printTree(){
    for(int i = 1;i<=n;i++){
        printf("cur: %d\n",i);
        printf("   chain Ind: %d , tree Pos: %d\n",chainInd[i],posInTree[i]);
        printf("   depth: %d , tree Array: %d\n",depth[i],treeArray[i]);
        printf("   belongs to %d subSize: %d\n",chainHead[chainInd[i]],subSize[i]);
    }
}

int main() {
    int nCase,Q,typ,u,v;
    long long val;
    scanf("%d ", &nCase);
    for (int cs = 1; cs <= nCase; cs++) {
        scanf("%d", &n);
        chainHead[0] = -1;
        for (int i = 1; i <= n; i++) {
            scanf("%lld",&Cost[i]);
            Graph[i].clear();
            chainHead[i] = -1;
            for (int j = 0; j < LN; j++) {
                sparseTable[i][j] = -1;
            }
        }
        for (int i = 0; i < n - 1; i++) {
            scanf("%d %d", &u, &v);
            u++,v++;
            Graph[u].push_back(v);
            Graph[v].push_back(u);
        }
        Node = 1;
        chainNo = 0;
        DFS(root, -1);
        HLD(root, -1);
        Node--;
        //printTree();
        build_tree(1, 1, Node);
        build_Sparse_Table();
        printf("Case %d:\n",cs);
        scanf("%d",&Q);

        for(int i = 0;i<Q;i++){
            scanf("%d",&typ);
            if (typ == 0) {
                scanf("%d %d",&u,&v);
                u++,v++;
                long long ans = query(u,v);
                printf("%lld\n", ans);
            } else {
                scanf("%d %lld",&u,&val);
                u++;
                updateNode(u, val);
            }
        }
    }
}