HackerEarth Mojtabas Trees and Arpas Queries March HourStorm solution

In this HackerEarth Mojtabas Trees and Arpas Queries March HourStorm problem solution Mojtaba has two trees, each of them has n vertices. Arpa has q queries, each in type v, u, x, y. Let s be the set of vertices in the path from v to u in the first tree and p be the set of vertices in the path from x to y in the second tree. Mojtaba has to calculate the size of s intersects p for each query. Help him!

HackerEarth Mojtabas Trees and Arpas Queries March HourStorm problem solution.

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

const int maxN = 300 * 1000 + 100; 
const int maxL = 20; 

typedef pair<int,int> pii; 

vector<int> c[2][maxN];
vector<pii> que[maxN];

int st[maxN], en[maxN];
int seg[4*maxN]; 

int h[2][maxN], par[2][maxN][maxL];

void dfs_lca(int t, int u) { 
    for(int k = 1; k < maxL; k++)
        par[t][u][k] = par[t][par[t][u][k-1]][k-1];

    for( auto x : c[t][u] )
        if( x != par[t][u][0] ) { 
            par[t][x][0] = u; 
            h[t][x] = h[t][u] + 1;
            dfs_lca(t, x);
        }
}

void dfs_time(int u, int p) {
    static int ind = 0;
    st[u] = ind++;
    for( auto x : c[0][u] ) 
        if( x != p ) 
            dfs_time(x, u);
    en[u] = ind;
}

int get_lca(int t, int u, int v) { 
    if( h[t][u] < h[t][v] ) swap(u, v);
    
    int diff = h[t][u] - h[t][v];
    for(int k = 0; k < maxL; k++)
        if( (diff>>k) & 1 ) 
            u = par[t][u][k];

    if( u == v ) return u; 

    for(int k = maxL - 1; k >= 0; k--)
        if( par[t][u][k] != par[t][v][k] ) { 
            u = par[t][u][k];
            v = par[t][v][k];
        }

    return par[t][u][0];
}

void query(int i, int u, int v, int x, int y) { 
    que[x].push_back( pii(i, u) );
    que[x].push_back( pii(-i, v) );
    que[y].push_back( pii(-i, u) );
    que[y].push_back( pii(i, v) );
}

int n;
int ans[maxN];

void seg_add(int ql, int qr, int qv, int xl=0, int xr=n, int ind=1) { 
    if( xr <= ql || qr <= xl ) return;
    if( ql <= xl && xr <= qr ) {
        seg[ind] += qv;
        return;
    }
    int xm = (xl+xr)/2;
    seg_add(ql, qr, qv, xl, xm, ind * 2);
    seg_add(ql, qr, qv, xm, xr, ind * 2 + 1);
}

int seg_get(int qp, int xl=0, int xr=n, int ind=1) { 
    if( xr - xl == 1 ) 
        return seg[ind];

    int xm = (xl+xr)/2;
    if( qp < xm ) 
        return seg[ind] + seg_get(qp, xl, xm, ind*2);
    return seg[ind] + seg_get(qp, xm, xr, ind*2+1);
}

void dfs_solve(int u, int p) {
    seg_add(st[u], en[u], 1);


    for(auto q: que[u]) { 
        int id = abs(q.first);
        int v = seg_get(st[q.second]);
        if( q.first < 0 ) 
            ans[id] -= v;
        else
            ans[id] += v;
    }

    for( auto x : c[1][u] ) 
        if( x != p ) 
            dfs_solve(x, u);

    seg_add(st[u], en[u], -1);
}

int main() { 
    ios::sync_with_stdio(false);
    cin.tie(0);

    int q;
    cin >> n >> q; 

    for(int t = 0; t < 2; t++) {
        c[t][0].push_back(1);
        for(int i = 0; i + 1 < n; i++) { 
            int u, v;
            cin >> u >> v;
            c[t][u].push_back(v);
            c[t][v].push_back(u);
        }
    }

    n++;

    dfs_lca(0, 0);
    dfs_lca(1, 0);

    dfs_time(0, -1);

    for(int i = 1; i <= q; i++) {
        int u, v, x, y;
        cin >> u >> v >> x >> y;

        int w = get_lca(0, u, v);
        int z = get_lca(1, x, y);


        query(i, u, par[0][w][0], x, par[1][z][0]);
        query(i, v, w, x, par[1][z][0]);
        query(i, v, w, y, z);
        query(i, u, par[0][w][0], y, z);
    }

    dfs_solve(0, -1);

    for(int i = 1; i <= q; i++)
        cout << ans[i] << 'n';

    return 0;
}

Second solution

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

const int maxn = 3e5 + 17, lg = 19;

struct Q{
    int x, y, i, z;
};
int n, q, par[2][lg][maxn], st[maxn], ft[maxn], h[2][maxn], ans[maxn], iman[maxn];
vector<int> g[2][maxn];
vector<Q> assign[maxn];
void make_par(int id, int v = 0){
    for(auto u : g[id][v])
        if(u != par[id][0][v]){
            par[id][0][u] = v;
            h[id][u] = h[id][v] + 1;
            make_par(id, u);
        }
}
void get_st(int v = 0){
    static int time = 0;
    st[v] = time++;
    for(auto u : g[1][v])
        if(u != par[1][0][v])
            get_st(u);
    ft[v] = time;
}
int lca(int id, int v, int u){
    if(h[id][v] > h[id][u])
        swap(v, u);
    for(int i = 0; i < lg; i++)
        if(h[id][u] - h[id][v] >> i & 1)
            u = par[id][i][u];
    for(int i = lg - 1; i >= 0; i--)
        if(par[id][i][v] != par[id][i][u])
            v = par[id][i][v], u = par[id][i][u];
    return v == u ? v : par[id][0][v];
}
int hamid(int p){
    int ans = 0;
    for(p++; p; p ^= p & -p)  ans += iman[p];
    return ans;
}
void majid(int p, int v){
    for(p++; p < maxn; p += p & -p)  iman[p] += v;
}
void majid(int l, int r, int v){
    majid(l, v), majid(r, -v);
}
void dfs(int v = 0){
    majid(st[v], ft[v], +1);
    for(auto q : assign[v]){
        int p = lca(1, q.x, q.y);
        ans[q.i] += q.z * (hamid(st[q.x]) + hamid(st[q.y]) - hamid(st[p]) - (p ? hamid(st[ par[1][0][p] ]) : 0));
    }
    for(auto u : g[0][v])
        if(u != par[0][0][v])
            dfs(u);
    majid(st[v], ft[v], -1);
}
int main(){
    ios::sync_with_stdio(0), cin.tie(0);
    cin >> n >> q;
    for(int k = 0; k < 2; k++)
        for(int i = 1; i < n; i++){
            int v, u;
            cin >> v >> u;
            v--, u--;
            g[k][v].push_back(u);
            g[k][u].push_back(v);
        }
    for(int j = 0; j < 2; j++){
        make_par(j);
        for(int k = 1; k < lg; k++)
            for(int v = 0; v < n; v++)
                par[j][k][v] = par[j][k - 1][ par[j][k - 1][v] ];
    }
    for(int i = 0; i < q; i++){
        int v, u, x, y, p;
        cin >> v >> u >> x >> y;
        v--, u--, x--, y--;
        p = lca(0, v, u);
        assign[v].push_back({x, y, i, +1});
        assign[u].push_back({x, y, i, +1});
        assign[p].push_back({x, y, i, -1});
        if(p)
            assign[ par[0][0][p] ].push_back({x, y, i, -1});
    }
    get_st();
    dfs();
    for(int i = 0; i < q; i++)
        cout << ans[i] << 'n';
}

coding problems solutions