In this HackerEarth Smart City problem solution, You live in a city that contains N houses. Each house is reachable from every other house through a unique path, which implies that there are N - 1 roads in the city. Each of the N - 1 road is distinct and has a length denoted by an array L. You want to reach directly from house i to house j (1 <= i, j <= N and i != j) by traversing not more than X units of road length. To reach directly from a house i to house j, you can select exactly one road of length less than or equal to X and destroy that road completely and construct a road between the two houses i and j but the condition is that the every house should be reachable from every other house.

Your task is to solve Q queries (each query is independent from each other ) of the following form:
A B X

where (A, B) denotes the pair of the houses that you want to connect. For each query, determine the number of ways in which you can select an ordered pair of houses (G,H) such that the following condidtions hold:
  1. There should be a road between the houses G and H
  2. The length of the road between the houses G and H is less than or equal to X
  3. After removing the road between G and H and adding a road between A and B, the houses G and H should still be connected through a path in the city.


HackerEarth Smart City problem solution


HackerEarth Smart City problem solution.

#include <bits/stdc++.h>
using namespace std;
#define logN 20
int n,q;
vector<pair<int,int> > v[100005];
vector<pair<int,int> > que1[100005],que2[100005];
int f[100005][21],depth[100005],tree[4*1000005],ans[100005];
bool vis[100005];
void dfs(int s)
{
    vis[s]=true;
    for (int i = 1; i<logN; i++)
        f[s][i] = f[f[s][i - 1]][i - 1];
    for(int i=0;i< v[s].size();i++){
        int node= v[s][i].first;
        if(vis[node]==false){
            f[node][0]=s;
            depth[node]=depth[s]+1;
            dfs(node);
        }
    }
}
int lca (int u, int v) {
    if (depth[u]<depth[v])
        swap(u, v);
    for (int i=logN-1;i >= 0;i--)
        if (depth[f[u][i]] >= depth[v]) {
            u = f[u][i];
        }
    if (u == v)
        return u;
    for (int i=logN -1;i>=0;i--)
        if (f[u][i] != f[v][i]) {
            u = f[u][i];
            v = f[v][i];
        }
    return f[u][0];
}
void update(int node,int start,int end,int idx,int val)
{
    if(start == end)
    {
        tree[node]+=val;
    }
    else
    {
        int mid = (start + end) / 2;
        if(start <= idx and idx <= mid)
        {
            update(2*node, start, mid, idx, val);
        }
        else
        {
            update(2*node+1, mid+1, end, idx, val);
        }
        tree[node]=(tree[2*node]+tree[2*node+1]);
    }
}
int query(int node,int start,int end,int l,int r)
{
    if(r < start || end < l)
    {
        return 0;
    }
    if(l <= start && end <= r)
    {
        return tree[node];
    }
    int mid = (start + end) / 2;
    int p1 = query(2*node, start, mid, l, r);
    int p2 = query(2*node+1, mid+1, end, l, r);
    return (p1 + p2);
}
void dfs1(int s,int len)
{
    vis[s]=true;
    if(len!=0){
        update(1,1,1000000,len,1);
    }
    for(auto i:que1[s]){
        ans[i.second]+= query(1,1,1000000,1,i.first);
    }
    for(auto i:que2[s]){
        ans[i.second]-= 2*query(1,1,1000000,1,i.first);
    }
    for(int i=0;i< v[s].size();i++){
        int node= v[s][i].first;
        if(vis[node]==false){
            dfs1(node,v[s][i].second);
        }
    }
    update(1,1,1000000,len,-1);
}
int main()
{
    //freopen("test1.txt","r",stdin);
    //freopen("file1.txt","w",stdout);
    cin>>n>>q;
    for(int i=1;i<=n-1;i++){
        int x,y,l;
        cin>>x>>y>>l;
        v[x].push_back({y,l});
        v[y].push_back({x,l});
    }
    depth[1]=1;
    dfs(1);
    memset(vis,false,sizeof(vis));
    for(int i=1;i<=q;i++)
    {
        int a,b,x;
        cin>>a>>b>>x;
        que1[a].push_back({x,i});
        que1[b].push_back({x,i});
        que2[lca(a,b)].push_back({x,i});
    }
    //cout<<"yes"<<endl;
    dfs1(1,0);
    for(int i=1;i<=q;i++){
        cout<<ans[i]<<endl;
    }
}

Second solution

#include <bits/stdc++.h>

using namespace std;

#define prev kek

const int M = 2e5 + 239;
const int N = (2 * M);
const int L = 19;
const int X = 1e6 + 239;

vector<pair<int, int> > v[M];
vector<int> ask[M], ans[M];
int n, q, a[M], b[M], lc[M], p1[M], p2[M], p3[M];
pair<int, int> dp[L][N];
vector<int> et;
int gl[M], first[N], how[N], k;

inline void dfs_lca(int p, int pr, int f)
{
    first[p] = et.size();
    et.push_back(p);
    gl[p] = f;
    for (pair<int, int> t : v[p])
        if (t.first != pr)
        {
            dfs_lca(t.first, p, f + 1);
            et.push_back(p);
        }
}

inline int lca(int s, int f)
{
    if (first[s] > first[f]) swap(s, f);
    s = first[s];
    f = first[f];
    int len = f - s + 1;
    int e = how[len];
    pair<int, int> ans = min(dp[e][s], dp[e][f + 1 - (1 << e)]);
    return ans.second;
}

inline void init_lca()
{
    dfs_lca(0, -1, 0);
    k = et.size();
    for (int i = 0; i < k; i++) dp[0][i] = make_pair(gl[et[i]], et[i]);
    for (int i = 1; i < L; i++)
        for (int j = 0; j < k; j++)
        {
            int st = (1 << (i - 1));
            if (st + j >= k)
            {
                dp[i][j] = dp[i - 1][j];
                continue;
            }
            dp[i][j] = min(dp[i - 1][j], dp[i - 1][st + j]);
        }
    how[1] = 0;
    for (int i = 2; i <= k; i++)
    {
        how[i] = how[i - 1];
        if ((1 << (how[i] + 1)) == i)
            how[i]++;
    }
}

int tree[X];

inline void upd(int i, int x)
{
    for (; i < X; i |= i + 1) tree[i] += x;
}

inline int gett(int i)
{
    int ans = 0;
    for (; i >= 0; i = (i & (i + 1)) - 1) ans += tree[i];
    return ans;
}

inline void dfs(int p, int pr)
{
    for (int i : ask[p])
        ans[p].push_back(gett(i));
    for (pair<int, int> t : v[p])
        if (t.first != pr)
        {
            upd(t.second, 1);
            dfs(t.first, p);
            upd(t.second, -1);
        }
}

int main()
{
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n >> q;
    for (int i = 0; i < n - 1; i++)
    {
        int s, f, c;
        cin >> s >> f >> c;
        s--, f--;
        v[s].push_back(make_pair(f, c));
        v[f].push_back(make_pair(s, c));
    }
    init_lca();
    for (int i = 0; i < q; i++)
    {
        int l;
        cin >> a[i] >> b[i] >> l;
        a[i]--, b[i]--;
        p1[i] = ask[a[i]].size();
        p2[i] = ask[b[i]].size();
        lc[i] = lca(a[i], b[i]);
        p3[i] = ask[lc[i]].size();
        ask[a[i]].push_back(l);
        ask[b[i]].push_back(l);
        ask[lc[i]].push_back(l);
    }
    dfs(0, -1);
    for (int i = 0; i < q; i++)
        cout << ans[a[i]][p1[i]] + ans[b[i]][p2[i]] - 2 * ans[lc[i]][p3[i]] << "\n";
    return 0;
}