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";
    }
}