[Codeforces] F. Tree Query

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : F. Tree Query
                    Hello 2015 (Div 1)
Source            : Codeforces
Category          : Graph Theory
Algorithm         : Centroid Decomposition
Verdict           : Accepted 

#include "bits/stdc++.h"
 
using namespace std;
 
static const int maxn = 1e5 + 5;
static const int logn = 18;
 
#define All(a)   a.begin(), a.end()
using pii = pair <long long, int>;
#define pb       push_back
 
struct node
{
    int v;
    long long w;
    node(int v = 0, long long w = 0) : v(v), w(w) {}
};
 
vector <node> graph[maxn];
vector <int> centroidtree[maxn];
int father[maxn][logn], depth[maxn];
long long fdis[maxn];
int subtreeSize[maxn], numnode, par[maxn], root_centroid_tree;
bool is_centroid[maxn];
int tnode, tquery, id;
 
void dfs(int u = 1, int p = -1)
{
    for (int i = 1; i < logn; i++)
        father[u][i] = father[ father[u][i-1] ][i-1];
    for (auto &it : graph[u])
    {
        int v = it.v;
        long long w = it.w;
        if (v == p) continue;
        depth[v] = depth[u] + 1;
        fdis[v] = fdis[u] + w;
        father[v][0] = u;
        dfs(v, u);
    }
}
 
int findlca(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];
}
 
long long dist(int u, int v)
{
    int lca = findlca(u, v);
    return fdis[u] + fdis[v] - 2*fdis[lca];
}
 
void dfs2(int u, int p = -1)
{
    subtreeSize[u] = 1;
    for (auto &it : graph[u])
    {
        int v = it.v;
        if (v == p || is_centroid[v]) continue;
        dfs2(v, u);
        subtreeSize[u] += subtreeSize[v];
    }
}
 
int getCentroid(int u, int p = -1)
{
    for (auto &it : graph[u])
    {
        int v = it.v;
        if (v == p || is_centroid[v]) continue;
        if (subtreeSize[v] > (numnode/2))
            return getCentroid(v, u);
    }
    return u;
}
 
int decompose(int u)
{
    dfs2(u);
    numnode = subtreeSize[u];
    int p = getCentroid(u);
    if (!root_centroid_tree) root_centroid_tree = p;
    is_centroid[p] = 1;
    for (auto &it : graph[p])
    {
        int v = it.v;
        if (is_centroid[v]) continue;
        int q = decompose(v);
        centroidtree[p].push_back(q);
        centroidtree[q].push_back(p);
    }
    return p;
}
 
void dfs4(int u, int p = -1)
{
    par[u] = p;
    for (int v : centroidtree[u])
    {
        if (v == p) continue;
        dfs4(v, u);
    }
}
 
// Problem F :
// It can also be solved by answering the queries online (using centroid decomposition).
// We decompose the given tree and for each node in the centroid tree,
// we maintain the following two vectors :
//
// ans[u] = sorted vector of d(x,u) for all x in the subtree of u in centroid tree,
//          where d(x,u) represents the distance between node x and u in the original tree.
// cntbn[u] = sorted vector of d(x,par[u]) for all x in the subtree of u in centroid tree,
//            where d(x,par[u]) represents the distance between node x and par[u] in the original tree.
//           (par[u] is the parent of node u in centroid tree).
//
// Since, there are at-most logN levels in centroid tree, and at each level we are storing O(N)
// values of ans/cntbn, total memory required would be O(NlogN). Also, these vectors can easily be
// build in O(NlogN) or O(Nlog^2N) time by moving up node u (for all u) in centroid tree and adding
// the corresponding values in all the ancestors of u.
//
// Given this information, each query would now reduce to
 
vector <long long> ans[maxn], cntbn[maxn];
 
void build_ans()
{
    for (int i = 1; i <= tnode; i++)
    {
        int u = i;
        while (1)
        {
            ans[u].pb(dist(u, i));
            if (par[u] == -1) break;
            u = par[u];
        }
    }
    for (int i = 1; i <= tnode; i++)
        sort(All(ans[i]));
}
 
void build_cntbn()
{
    for (int i = 1; i <= tnode; i++)
    {
        int u = i;
        while (1)
        {
            if (par[u] == -1) break;
            cntbn[u].pb(dist(par[u], i));
            u = par[u];
        }
    }
    for (int i = 1; i <= tnode; i++)
        sort(All(cntbn[i]));
}
 
int query(int u, long long l)
{
    int counter = upper_bound(All(ans[u]), l) - ans[u].begin();
    int src = u;
    while (1)
    {
        if (par[u] == -1) break;
        long long d = l - dist(src, par[u]);
        int jog = upper_bound(All(ans[ par[u] ]), d) - ans[ par[u] ].begin();
        int bad = upper_bound(All(cntbn[u]), d) - cntbn[u].begin();
        counter += jog - bad;
        u = par[u];
    }
    return counter;
}
 
 
int main()
{
    //freopen("in.txt", "r", stdin);
 
    scanf("%d %d", &tnode, &tquery);
    for (int e = 1; e < tnode; e++)
    {
        int u, v;
        long long w;
        scanf("%d %d %lld", &u, &v, &w);
        graph[u].push_back(node(v, w));
        graph[v].push_back(node(u, w));
    }
    depth[1] = 1;
    dfs(1);
    decompose(1);
    dfs4(root_centroid_tree);
    build_ans();
    build_cntbn();
    while (tquery--)
    {
        int v;
        long long l;
        cin >> v >> l;
        printf("%d\n", query(v, l));
    }
}