HackerEarth Subtree swap problem solution

In this HackerEarth Subtree swap problem solution, You are given an undirected tree with N nodes rooted at node 1. Every node has a value A[i] assigned to it. You are required to answer Q queries of the following type:

UVX

Select a subtree with U as the root node and subtree with V as the root node and swap their positions. That is, detach both the subtrees and swap their positions.

If node U is an ancestor of node V or node V is an ancestor of node U, the above Swap operation is not performed.

Find the sum of node values present in the subtree rooted at node X.

If the Swap operation is performed, then redo this operation. That is, swap the subtree with U as the root node and subtree with V as the root node.

Determine the required answer for Q queries.

HackerEarth Subtree swap problem solution.

#include<bits/stdc++.h>
using namespace std;
int tim = 0;
int in_t[100001];
int out_t[100001];
long long dp[100001];
int dep[100001];
int is_anc(int u, int v) // check if 'u' is ancestor of 'v'
{
    return (in_t[u] <= in_t[v] && in_t[v] < out_t[u]);
}
void dfs(vector<int> tree[], int v, int p, vector<int> &A)
{
    dp[v] = A[v];
    dep[v] = 1 + dep[p];
    in_t[v] = ++tim;
    for(auto j : tree[v])
        if(j != p)
            dfs(tree, j, v, A), dp[v]+=dp[j];
    out_t[v] = ++tim;
}
vector<long long> solve (int N, vector<int> A, vector<vector<int> > edges, int Q, vector<vector<int> > query) {
    // Write your code here
    vector<int> tree[N+1];
    int i;
    for(i=0;i<N-1;i++)
    {
        assert(edges[i][0]!=edges[i][1]);
        tree[edges[i][0]].push_back(edges[i][1]);
        tree[edges[i][1]].push_back(edges[i][0]);
    }
    tim = 0;
    reverse(A.begin(), A.end());
    A.push_back(0);
    reverse(A.begin(), A.end());
    dfs(tree, 1, 0, A);
    for(int i=2;i<=N;i++) assert(dep[i]);
    vector<long long> ans;
    for(i = 0; i < Q ; i++)
    {
        int u = query[i][0], v = query[i][1], x = query[i][2];
        if(is_anc(u, v) || is_anc(v, u))
        {
            ans.push_back(dp[x]);
            continue;
        }
        assert(1 <= u && u <= N);
        assert(1 <= v && v <= N);
        assert(1 <= x && x <= N);
        if(x == u || x == v)
        {
            ans.push_back(dp[x]);
            continue;
        }
        int f1 = is_anc(x, u);
        int f2 = is_anc(x, v);
        if(f1 == 1 && f2 == 1)
            ans.push_back(dp[x]);
        else if(f1 == 1)
            ans.push_back(dp[x] - dp[u] + dp[v]);
        else if(f2 == 1)
            ans.push_back(dp[x] - dp[v] + dp[u]);
        else
            ans.push_back(dp[x]);
    }
    return ans;
}

int main() {

    ios::sync_with_stdio(0);
    cin.tie(0);
    int T;
    cin >> T;
    assert(1 <= T && T <= 10);
    for(int t_i = 0; t_i < T; t_i++)
    {
        int N;
        cin >> N;
        assert(1 <= N && N <= 1e5);
        vector<int> A(N);
        for(int i_A = 0; i_A < N; i_A++)
        {
            cin >> A[i_A];
            assert(1 <= A[i_A] && A[i_A] <= 1e9);
        }
        vector<vector<int> > edges(N-1, vector<int>(2));
        for (int i_edges = 0; i_edges < N-1; i_edges++)
        {
            for(int j_edges = 0; j_edges < 2; j_edges++)
            {
                cin >> edges[i_edges][j_edges];
                assert(1<=edges[i_edges][j_edges] and edges[i_edges][j_edges]<=N);;
            }
        }
        int Q;
        cin >> Q;
        assert(1 <= Q && Q <= 1e5);
        vector<vector<int> > query(Q, vector<int>(3));
        for (int i_query = 0; i_query < Q; i_query++)
        {
            for(int j_query = 0; j_query < 3; j_query++)
            {
                cin >> query[i_query][j_query];
            }
        }

        vector<long long> out_;
        out_ = solve(N, A, edges, Q, query);
        cout << out_[0];
        for(int i_out_ = 1; i_out_ < out_.size(); i_out_++)
        {
            cout << " " << out_[i_out_];
        }
        cout << "\n";
    }
}