[UVa] 12238 - Ant Colony

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : 12238 - Ant Colony
Source            : UVa Online Judge
Category          : Data Structure, Tree
Algorithm         : MO's Algorithm on Tree
Verdict           : Accepted

#include <bits/stdc++.h>
 
 
using namespace std;
 
typedef long long          ll;
 
static const int MAXN = 2e5+15;
static const int LOGN = 19;
static const int BLOCK = 450;
 
struct node
{
    int v;
    ll w;
};
 
vector <node> graph[MAXN];
 
int nodeList[MAXN], ST[MAXN], ET[MAXN], dfsTime;
ll weight[MAXN];
 
int pointer, depth[MAXN], father[MAXN][LOGN];
 
void dfs(int u, int p)
{
    dfsTime++;
    pointer++;
    ST[u] = dfsTime;
    nodeList[pointer] = u;
    for (int i = 1; i < LOGN; i++)
    {
        father[u][i] = father[father[u][i-1]][i-1];
    }
    for (auto e : graph[u])
    {
        int v = e.v;
        ll w = e.w;
        if (v == p) continue;
        depth[v] = depth[u] + 1;
        father[v][0] = u;
        weight[v] = w;
        dfs(v, u);
    }
    dfsTime++;
    pointer++;
    ET[u] = dfsTime;
    nodeList[pointer] = u;
}
 
int LCA(int u, int v)
{
    if (depth[u] < depth[v]) swap(u, v);
    for (int i = LOGN-1; i >= 0; i--)
    {
        if (depth[father[u][i]] >= depth[v])
        {
            u = father[u][i];
        }
    }
    if (u == v) return u;
    for (int i = LOGN-1; i >= 0; i--)
    {
        if (father[u][i] != father[v][i])
        {
            u = father[u][i];
            v = father[v][i];
        }
    }
    return father[u][0];
}
 
 
struct MO
{
    int l, r, id;
    MO() {}
    MO(int _l, int _r, int _id)
    {
        l = _l;
        r = _r;
        id = _id;
    }
    friend bool operator<(MO A, MO B)
    {
        int BA = A.l / BLOCK;
        int BB = B.l / BLOCK;
        if (BA == BB) return A.r < B.r;
        return BA < BB;
    }
} Q[MAXN];

int l, r;
ll ans, ANS[MAXN];
int frequencyNode[MAXN];
 
void Add(int pos)
{
    int n = nodeList[pos];
    ll w = weight[n];
    frequencyNode[n]++;
    if (frequencyNode[n] == 1) ans += w;
    else if (frequencyNode[n] == 2) ans -= w;
}
 
void Remove(int pos)
{
    int n = nodeList[pos];
    ll w = weight[n];
    frequencyNode[n]--; 
    if (frequencyNode[n] == 1) ans += w;
    else if (frequencyNode[n] == 0) ans -= w;
}
 
void CLEAR()
{
    ans = 0;
    dfsTime = 0;
    pointer = 0;
    for (int i = 0; i < MAXN; i++)
    {
        weight[i] = 0;
        depth[i] = 0;
        ST[i] = 0;
        ET[i] = 0;
        frequencyNode[i] = 0;
        nodeList[i] = 0;
        graph[i].clear();
        for (int j = 0; j < LOGN; j++) father[i][j] = 0;
    }
}
 
int main()
{
    //freopen("in.txt", "r", stdin);
 
    int tNode;
    while (scanf("%d", &tNode) == 1)
    {
        if (tNode == 0) break;
        CLEAR();
        for (int i = 1; i < tNode; i++)
        {
            int v;
            ll w;
            scanf("%d %lld", &v, &w);
            int u = i + 1;
            v++;
            graph[u].push_back({v, w});
            graph[v].push_back({u, w});
        }
        depth[1] = 1;
        dfs(1, -1);
        int tQuery;
        scanf("%d", &tQuery); 
        for (int i = 0; i < tQuery; i++)
        {
            int u, v;
            scanf("%d %d", &u, &v);
            u++;
            v++;
            if (ST[u] > ST[v]) swap(u, v);
            int lca = LCA(u, v);
            if (lca == u)
            {
                Q[i] = {ST[u]+1, ST[v], i};
            }
            else
            {
                Q[i] = {min(ET[u], ST[v]), max(ET[u], ST[v]), i};
            }
        }
        sort(Q, Q+tQuery);
        l = 2;
        r = 1;
        for (int i = 0; i < tQuery; i++)
        {
            int L = Q[i].l;
            int R = Q[i].r;
            int ID = Q[i].id;
            while (l > L)      Add(--l);
            while (r < R)      Add(++r);
            while (l < L)      Remove(l++);
            while (r > R)      Remove(r--);
            ANS[ID] = ans;
        }
        for (int i = 0; i < tQuery; i++)
        {
            printf("%lld", ANS[i]);
            printf(i+1==tQuery ? "\n" : " ");
        }
    }
}