Spoj QTREE2 Solution

on

#include<bits/stdc++.h>
#define pii pair<int, int>
#define mx 10005
#define fast ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL);
using namespace std;

int level[mx], parent[mx], st[24][mx], dis[mx];
vector<int>a[mx];
map<pii, int>mp;

void dfs(int from, int u, int depth)
{
    level[u] = depth;
    parent[u] = from;
    st[0][u] = from;
    dis[u] = dis[from] + mp[pii(from, u)];

    for(int i=0; i<a[u].size(); i++)
    {
        int next = a[u][i];

        if(from == next)
            continue;

        dfs(u, next, depth+1);
    }
}

void initialize(int n)
{
    for(int k=1; k<=16; k++)
    {
        for(int i=0; i<n; i++)
        {
            if(st[k-1][i] != -1)
                st[k][i] = st[k-1][st[k-1][i]];
        }
    }
}

int lca(int p, int q)
{
    if(level[p] < level[q])
        swap(p, q);

    int log = 1;
    while(1)
    {
        int next = log+1;

        if((1<<next) > level[p])
            break;

        log++;
    }

    for(int i=log; i>=0; i--)
    {
        if(level[p]-(1<<i) >= level[q])
            p = st[i][p];
    }

    if(p == q)
        return p;

    for(int i=log; i>=0; i--)
    {
        if(st[i][p] != -1 && st[i][p] != st[i][q])
        {
            p = st[i][p];
            q = st[i][q];
        }
    }

    return parent[p];
}

int calculate(int x, int y, int k)
{
    int r = lca(x, y);
    int length = level[x] - level[r] + 1;

    if(k <= length)
    {
        k--;

        for(int i=0; (k>>i) != 0; i++)
        {
            if((k>>i)&1)
                x = st[i][x];
        }

        return x;
    }

    else
    {
        k = level[y] - level[r] - (k-length);

        for(int i=0; (k>>i) != 0; i++)
        {
            if((k>>i)&1)
                y = st[i][y];
        }

        return y;
    }
}

void clears()
{
    mp.clear();
    for(int i=0; i<24; i++)
    {
        for(int j=0; j<mx; j++)
            st[i][j] = -1;
    }

    for(int i=0;i<mx;i++)
        level[i] = parent[i] = 0, a[i].clear();
}

int main()
{
    fast

    int t;
    cin >> t;
    while(t--)
    {
        int n, e, x, y, value;
        cin >> n;

        for(int i=1; i<=n-1; i++)
        {
            cin >> x >> y >> value;
            x--, y--;
            a[x].push_back(y);
            mp[pii(x, y)] = value;
            mp[pii(y, x)] = value;
        }

        dfs(-1, 0, 0);
        initialize(n);

        string s;
        while(cin >> s)
        {
            if(s == "DONE")
                break;

            if(s == "DIST")
            {
                cin >> x >> y;
                x--, y--;

                int r = lca(x, y);

                int sum = abs(dis[r] - dis[x]) + abs(dis[r] - dis[y]);

                cout << sum << "\n";
            }
            else
            {
                int k;
                cin >> x >> y >> k;
                x--, y--;

                int r = calculate(x, y, k);
                r++;
                cout << r << "\n";
            }
        }

        clears();
    }

    return 0;
}

Comments

Post a Comment