In this HackerEarth Ambiguous Tree problem solution For his birthday, Jeffrey received a tree with N vertices, numbered from 1 to N. Disappointed by the fact that there is only one simple path between any two vertices, Jeffrey decided to remedy this by introducing an extra edge in the tree. Note that this may result in a self loop or multiple edge. After introducing an extra edge, Jeffrey wants to know how many pairs of two vertices are ambiguous. Two vertices are ambiguous if there exists more than one simple path between them. Jeffrey has Q queries — he has decided Q possible locations for the extra edge. Please help him find out how many ambiguous pairs there in the tree if he added each of the Q edges to the tree. The tree will revert back to its original state after each query, so queries are mutually independent.


HackerEarth Ambiguous Tree problem solution


HackerEarth Ambiguous Tree problem solution.

#include <bits/stdc++.h>

using namespace std;

const int SZ=350;
int N, Q;
vector<int> adj[100001];
long long C2[100001];
int level[100001];
int sz[100001];
int pos[100001];
int preorder[200001], npreorder;
int P[18][100001];
bool cycle[100001];
long long ans[100001];
long long curans;

struct info
{
    int l, r, idx;
    bool operator< (const info& other) const
    {
        if(pos[l]/SZ!=pos[other.l]/SZ)
            return pos[l]/SZ<pos[other.l]/SZ;
        return pos[r]<pos[other.r];
    }
} queries[100001];

void dfs(int u, int p)
{
    sz[u]=1;
    pos[u]=npreorder;
    preorder[npreorder++]=u;
    P[0][u]=p;
    for(int i=1; i<=17; i++)
        P[i][u]=P[i-1][P[i-1][u]];
    for(auto& v: adj[u]) if(v!=p)
    {
        level[v]=level[u]+1;
        dfs(v, u);
        preorder[npreorder++]=-u;
        sz[u]+=sz[v];
    }
}

int lca(int u, int v)
{
    if(level[u]<level[v])
        swap(u, v);
    int uu=u, vv=v;
    for(int i=17; i>=0; i--) if(level[P[i][uu]]>=level[vv])
        uu=P[i][uu];
    for(int i=17; i>=0; i--) if(P[i][uu]!=P[i][vv])
        uu=P[i][uu], vv=P[i][vv];
    if(uu==vv)
        return uu;
    return P[0][uu];
}

inline void move_add(int u, int v)
{
    cycle[v]=true;
    curans-=C2[sz[u]];
    sz[u]-=sz[v];
    curans+=C2[sz[u]];
    curans+=C2[sz[v]];
}

inline void move_rem(int u, int v)
{
    cycle[v]=false;
    curans-=C2[sz[v]];
    curans-=C2[sz[u]];
    sz[u]+=sz[v];
    curans+=C2[sz[u]];
}

inline void move_path_up(int u, int v)
{
    if(u==v)
        return;
    if(cycle[P[0][u]])
        move_rem(P[0][u], u);
    else
        move_add(u, P[0][u]);
    move_path_up(P[0][u], v);
}

inline void move_path_down(int u, int v)
{
    if(u==v)
        return;
    move_path_down(P[0][u], v);
    if(cycle[u])
        move_rem(u, P[0][u]);
    else
        move_add(P[0][u], u);
}

int main()
{
    scanf("%d", &N);
    for(int i=1; i<=N; i++)
        C2[i]=1LL*i*(i-1)/2;
    int a, b;
    for(int i=1; i<N; i++)
    {
        scanf("%d%d", &a, &b);
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    dfs(1, 1);
    scanf("%d", &Q);
    for(int i=0; i<Q; i++)
    {
        scanf("%d%d", &a, &b);
        if(a>b)
            swap(a, b);
        queries[i].l=a;
        queries[i].r=b;
        queries[i].idx=i;
    }
    sort(queries, queries+Q);
    curans=1LL*N*(N-1)/2;
    int u=1, v=1;
    cycle[1]=true;
    for(int i=0; i<Q; i++)
    {
        int x=lca(queries[i].l, u);
        move_path_up(u, x);
        move_path_down(queries[i].l, x);
        x=lca(queries[i].r, v);
        move_path_up(v, x);
        move_path_down(queries[i].r, x);
        ans[queries[i].idx]=C2[N]-curans;
        u=queries[i].l;
        v=queries[i].r;
    }
    for(int i=0; i<Q; i++)
        printf("%lld\n", ans[i]);
    return 0;
}

Second solution

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

#define jjs(i, s, x) for (int i = (s); i < int(x); i++)
#define jjl(i, x) jjs(i, 0, x)
#define ji(x) jjl(i, x)
#define jj(x) jjl(j, x)
#define jk(x) jjl(k, x)
#define jij(a, b) ji(a) jj(b)
#define ever ;;
#define foreach(x, C) for (auto& x : (C))
#define INF ((int) 1e9+10)
#define LINF ((ll) 1e16)
#define pb push_back
#define mp make_pair
#define nrint(x) int x; rint(x);
#define nrlong(x) long long x; rint(x);
#define rfloat(x) scanf("%lf", &(x))

#ifndef ONLINE_JUDGE
const bool DEBUG = true;
#define Db(x...)   ({ if (DEBUG) { cout << "Debug["; DB, #x, ":", x, "]\n"; } })
template<class T> void Dbrng(T lo, T hi, string note = "", int w = 0) {
  if (DEBUG) {  
    cout << "Debug[ ";
    if (!note.empty()) cout << setw(3) << note << " : ";
    for (; lo != hi; ++lo) cout << setw(w) << *lo << " ";
    cout << "]" << endl;
  }
}
struct Debugger { template<class T> Debugger& operator ,
  (const T & v) { cout << " " << v << flush; return *this; } } DB;
#else
const bool DEBUG = false;
#define Db(x...)
#endif

#define rint readInteger
template<typename T>
bool readInteger(T& x)
{
    char c,r=0,n=0;
    x=0;
    for (ever)
    {
        c=getchar();
        if ((c<0) && (!r))
            return(0);
        else if ((c=='-') && (!r))
            n=1;
        else if ((c>='0') && (c<='9'))
            x=x*10+c-'0',r=1;
        else if (r)
            break;
    }
    if (n)
        x=-x;
    return(1);
}

const int MOD = 1000000007;
typedef long long ll;
typedef pair<int, int> pi;

const int MX = 1e5;
const int LOG = 20;

struct Answer
{
    ll pairs, sz;

    Answer operator ^ (const Answer& o) const;
    Answer& operator ^= (const Answer& o)
    {
        return *this = *this ^ o;
    }
};
Answer combine(Answer a, Answer b)
{
    return {b.sz * a.sz + a.pairs + b.pairs, a.sz + b.sz};
}
Answer Answer::operator ^ (const Answer& o) const
{
    return combine(*this, o);
}

struct node
{
    vector<int> edges;
    int depth;
    int subtreeSize;
    int parent;
    int jmp[LOG];
    Answer fol[LOG];
} n[MX];

void init1(int x, int parent)
{
    if (parent >= 0)
        n[x].depth = n[parent].depth + 1;
    n[x].parent = parent;
    n[x].subtreeSize = 1;
    for (int o : n[x].edges) if (o != parent)
    {
        init1(o, x);
        n[x].subtreeSize += n[o].subtreeSize;
    }
}

void init2(int x, int parent)
{
    n[x].jmp[0] = parent;
    if (n[x].parent >= 0)
        n[x].fol[0] = {0, n[parent].subtreeSize - n[x].subtreeSize};
    for (int i = 1; i < LOG; i++)
    {
        if (n[x].jmp[i-1] < 0)
            n[x].jmp[i] = -1;
        else
        {
            n[x].jmp[i] = n[n[x].jmp[i-1]].jmp[i-1];
            n[x].fol[i] = n[x].fol[i-1] ^ n[n[x].jmp[i-1]].fol[i-1];
        }
    }
    for (int o : n[x].edges) if (o != parent)
        init2(o, x);
}

int LCA(int a, int b)
{
    if (n[a].depth < n[b].depth)
        swap(a, b);
    for (int i = LOG-1; i >= 0; i--)
    {
        if (n[a].jmp[i] >= 0 && n[n[a].jmp[i]].depth >= n[b].depth)
            a = n[a].jmp[i];
    }
    if (a == b) return a;
    for (int i = LOG-1; i >= 0; i--)
    {
        if (n[a].jmp[i] != n[b].jmp[i])
        {
            a = n[a].jmp[i];
            b = n[b].jmp[i];
        }
    }
    return n[a].parent;
}

int follow(int a, int target, Answer& ans)
{
    assert(n[target].depth <= n[a].depth);
    if (a == target) return 0;
    ans ^= {0, n[a].subtreeSize};
    for (int i = LOG-1; i >= 0; i--)
    {
        if (n[a].jmp[i] >= 0 && n[n[a].jmp[i]].depth > n[target].depth)
        {
            ans ^= n[a].fol[i];
            a = n[a].jmp[i];
        }
    }
    return n[a].subtreeSize;
}

int N, Q;

int main()
{
    rint(N);
    assert(1 <= N && N <= MX);
    jk(N-1)
    {
        nrint(a); nrint(b);
        assert(1 <= a && a <= N);
        assert(1 <= b && b <= N);
        a--; b--;
        n[a].edges.pb(b);
        n[b].edges.pb(a);
    }
    init1(0, -1);
    init2(0, -1);
    rint(Q);
    assert(1 <= Q && Q <= MX);
    ji(Q)
    {
        nrint(a); nrint(b);
        assert(1 <= a && a <= N);
        assert(1 <= b && b <= N);
        a--; b--;
        int u = LCA(a, b);
        Answer ans = {0, 0};
        int v1 = follow(a, u, ans);
        int v2 = follow(b, u, ans);
        ans ^= {0, N - v1 - v2};
        printf("%lld\n", ans.pairs);
    }
}