HackerEarth Colorful Tree problem solution

YASH PAL,
In this HackerEarth Colorful Tree problem solution You are given a tree that contains N nodes, where every node i is colored with some color Ci.
The distance of a node V from a node U is defined as the number of edges along the simple path from the node U to the node V. Your task is to answer M queries of the following type:
  •  K C: Determine the distance of most distant node of color C from node K. If there is no node of color C in the tree, then print -1.
HackerEarth Colorful Tree problem solution

HackerEarth Colorful Tree problem solution.

#include <bits/stdc++.h>
using namespace std;
 
#define IOS ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#define endl "n"

const int N=5e5+5;
const int LG=21;

int n, k, q, tim=0, dist=0, node;
int col[N];
int parent[LG][N];
int tin[N], tout[N], level[N], vertices[2*N];
vector<int> g[N], contains[N], diam[N];
vector<pair<int, int> > tree[N];

void dfs(int k, int par, int lvl)
{
    tin[k]=++tim;
    parent[0][k]=par;
    level[k]=lvl;
    for(auto it:g[k])
    {
        if(it==par)
            continue;
        dfs(it, k, lvl+1);
    }
    tout[k]=tim;
}

void precompute()
{
    for(int i=1;i<LG;i++)
        for(int j=1;j<=n;j++)
            if(parent[i-1][j])
                parent[i][j]=parent[i-1][parent[i-1][j]];
}

int LCA(int u, int v)
{
    if(level[u]<level[v])
        swap(u,v);
    int diff=level[u]-level[v];
    for(int i=LG-1;i>=0;i--)
    {
        if((1<<i) & diff)
        {
            u=parent[i][u];
        }
    }
    if(u==v)
        return u;
    for(int i=LG-1;i>=0;i--)
    {
        if(parent[i][u] && parent[i][u]!=parent[i][v])
        {
            u=parent[i][u];
            v=parent[i][v];
        }
    }
    return parent[0][u];
}

int dist1(int u, int v)
{
    return level[u]+level[v]-2*level[LCA(u, v)];
}

bool isancestor(int u, int v) //Check if u is an ancestor of v
{
    return (tin[u]<=tin[v]) && (tout[v]<=tout[u]);
}

int dfs2(int k, int par, int dis)
{
    //cerr<<k<<endl;
    if(dis>dist)
    {
        dist=dis;
        node=k;
    }
    for(auto it:tree[k])
    {
        if(it.first==par)
            continue;
        dfs2(it.first, k, dis+it.second);
    }
}

int work(int color)
{
    sort(vertices+1, vertices+k+1, [](int a, int b)
    {
        return tin[a]<tin[b];
    });
    int idx=k;
    for(int i=1;i<idx;i++)
        vertices[++k]=LCA(vertices[i], vertices[i+1]);
    sort(vertices+1, vertices+k+1);
    k=unique(vertices+1, vertices+k+1) - vertices - 1;
    sort(vertices+1, vertices+k+1, [](int a, int b)
    {
        return tin[a]<tin[b];
    });
    stack<int> s;
    s.push(vertices[1]);
    for(int i=2;i<=k;i++)
    {
        while(!isancestor(s.top(), vertices[i]))
            s.pop();
        int u=s.top();
        int v=vertices[i];
        int w=dist1(u, v);
        tree[u].push_back({v, w});
        tree[v].push_back({u, w});
        s.push(vertices[i]);
    }
    dist=0;
    dfs2(vertices[1], vertices[1], 1);
    diam[color].push_back(node);
    dfs2(node, node, 1);
    diam[color].push_back(node);
    for(int i=1;i<=k;i++)
        tree[vertices[i]].clear();
}

int32_t main()
{
    IOS;
    cin>>n>>q;
    for(int i=1;i<=n;i++)
    {
        cin>>col[i];
        contains[col[i]].push_back(i);
    }
    for(int i=1;i<=n-1;i++)
    {
        int u, v;
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs(1, 0, 1);
    precompute();
    for(int i=1;i<=5e5;i++)
    {
        if(contains[i].size()<=2)
        {
            for(auto &it:contains[i])
                diam[i].push_back(it);
            continue;
        }
        k=0;
        for(auto &it:contains[i])
            vertices[++k]=it;
        work(i);
    }
    while(q--)
    {
        int k, c;
        cin>>k>>c;
        if(!contains[c].size())
            cout<<"-1"<<endl;
        else 
        {
            int ans=0;
            for(auto &it: diam[c])
                ans=max(ans, dist1(k, it));
            cout<<ans<<endl;
        }
    }
    return 0;
}

Second solution

#ifndef _GLIBCXX_NO_ASSERT
#include <cassert>
#endif
#include <cctype>
#include <cerrno>
#include <cfloat>
#include <ciso646>
#include <climits>
#include <clocale>
#include <cmath>
#include <csetjmp>
#include <csignal>
#include <cstdarg>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>

#if __cplusplus >= 201103L
#include <ccomplex>
#include <cfenv>
#include <cinttypes>
#include <cstdbool>
#include <cstdint>
#include <ctgmath>
#include <cwchar>
#include <cwctype>
#endif

// C++
#include <algorithm>
#include <bitset>
#include <complex>
#include <deque>
#include <exception>
#include <fstream>
#include <functional>
#include <iomanip>
#include <ios>
#include <iosfwd>
#include <iostream>
#include <istream>
#include <iterator>
#include <limits>
#include <list>
#include <locale>
#include <map>
#include <memory>
#include <new>
#include <numeric>
#include <ostream>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdexcept>
#include <streambuf>
#include <string>
#include <typeinfo>
#include <utility>
#include <valarray>
#include <vector>

#if __cplusplus >= 201103L
#include <array>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <forward_list>
#include <future>
#include <initializer_list>
#include <mutex>
#include <random>
#include <ratio>
#include <regex>
#include <scoped_allocator>
#include <system_error>
#include <thread>
#include <tuple>
#include <typeindex>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>
#endif

#define ll          long long
#define pb          push_back
#define mp          make_pair
#define pii         pair<int,int>
#define vi          vector<int>
#define all(a)      (a).begin(),(a).end()
#define F           first
#define S           second
#define sz(x)       (int)x.size()
#define hell        1000000007
#define endl        'n'
#define rep(i,a,b)  for(int i=a;i<b;i++)
using namespace std;

string to_string(string s) {
    return '"' + s + '"';
}

string to_string(const char* s) {
    return to_string((string) s);
}

string to_string(bool b) {
    return (b ? "true" : "false");
}

string to_string(char ch) {
    return string("'")+ch+string("'");
}

template <typename A, typename B>
string to_string(pair<A, B> p) {
    return "(" + to_string(p.first) + ", " + to_string(p.second) + ")";
}

template <class InputIterator>
string to_string (InputIterator first, InputIterator last) {
    bool start = true;
    string res = "{";
    while (first!=last) {
        if (!start) {
            res += ", ";
        }
        start = false;
        res += to_string(*first);
        ++first;
    }
    res += "}";
    return res;
}

template <typename A>
string to_string(A v) {
    bool first = true;
    string res = "{";
    for (const auto &x : v) {
        if (!first) {
            res += ", ";
        }
        first = false;
        res += to_string(x);
    }
    res += "}";
    return res;
}

void debug_out() { cerr << endl; }

template <typename Head, typename... Tail>
void debug_out(Head H, Tail... T) {
    cerr << " " << to_string(H);
    debug_out(T...);
}

template <typename A, typename B>
istream& operator>>(istream& input,pair<A,B>& x){
    input>>x.F>>x.S;
    return input;
}

template <typename A>
istream& operator>>(istream& input,vector<A>& x){
    for(auto& i:x)
        input>>i;
    return input;
}


#ifdef PRINTERS
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

long long readInt(long long l,long long r,char endd){
    long long x=0;
    int cnt=0;
    int fi=-1;
    bool is_neg=false;
    while(true){
        char g=getchar();
        if(g=='-'){
            assert(fi==-1);
            is_neg=true;
            continue;
        }
        if('0'<=g && g<='9'){
            x*=10;
            x+=g-'0';
            if(cnt==0){
                fi=g-'0';
            }
            cnt++;
            assert(fi!=0 || cnt==1);
            assert(fi!=0 || is_neg==false);

            assert(!(cnt>19 || ( cnt==19 && fi>1) ));
        } else if(g==endd){
            assert(cnt>0);
            if(is_neg){
                x= -x;
            }
            assert(l<=x && x<=r);
            return x;
        } else {
            debug(int(g));
            assert(false);
        }
    }
}
string readString(int l,int r,char endd){
    string ret="";
    int cnt=0;
    while(true){
        char g=getchar();
        if(g==endd){
            break;
        }
        else if(islower(g)){
            cnt++;
            ret+=g;
        }
        else{
            assert(false);
        }
    }
    assert(l<=cnt && cnt<=r);
    return ret;
}
long long readIntSp(long long l,long long r){
    return readInt(l,r,' ');
}
long long readIntLn(long long l,long long r){
    return readInt(l,r,'n');
}
string readStringLn(int l,int r){
    return readString(l,r,'n');
}
string readStringSp(int l,int r){
    return readString(l,r,' ');
}

vi colnode[500005];

int intime[500005];
int outtime[500005];
int height[500005];
vi adj[500005];
int dp[20][500005];
vector<pii> temptree[500005];
vi reqdnodes[500005];
int color[500005];


void dfs(int u,int p=0){
    static int clck = 1;
    intime[u] = clck;
    colnode[color[u]].emplace_back(u);
    height[u] = height[p]+1;
    dp[0][u]=p;
    clck++;
    for(auto i:adj[u]){
        if(i!=p)dfs(i,u);
    }
    outtime[u]=clck;
    clck++;
}

int lca(int u,int v){
    if(height[u]>height[v])swap(u,v);
    for(int i=19;i>=0;i--){
        if(height[v]-(1<<i)>=height[u])v=dp[i][v];
    }
    if(u==v)return u;
    for(int i=19;i>=0;i--){
        if(dp[i][u]!=dp[i][v])u=dp[i][u],v=dp[i][v];
    }
    return dp[0][u];
}

struct diameter
{
    pii maxdep1,maxdep2;
    pair<int,pii> best_res;
    diameter(int u){
        maxdep1.S=u;
        maxdep2.S=u;
        best_res.S={u,u};
    }
};

diameter get_diameter(int u,int p){
    diameter res(u);
    for(auto i:temptree[u]){
        if(i.F==p)continue;
        auto new_res = get_diameter(i.F,u);
        if(new_res.maxdep1.F+i.S>res.maxdep1.F){
            res.maxdep2=res.maxdep1;
            res.maxdep1=mp(new_res.maxdep1.F+i.S,new_res.maxdep1.S);
        }
        else if(new_res.maxdep1.F+i.S>res.maxdep2.F){
            res.maxdep2=mp(new_res.maxdep1.F+i.S,new_res.maxdep1.S);
        }
        res.best_res=max(res.best_res,new_res.best_res);
    }
    res.best_res=max(res.best_res,mp(res.maxdep1.F+res.maxdep2.F,mp(res.maxdep1.S,res.maxdep2.S)));
    return res;
}

void solve(){
    auto comp = [](int a,int b){return intime[a]<intime[b];};
    int N,M;
    N = readIntSp(1,500000);
    M = readIntLn(1,500000);
    rep(i,1,N+1){
        int col;
        if(i==N) col = readIntLn(1,500000);
        else col = readIntSp(1,500000);
        color[i] = col;
    }
    rep(i,1,N){
        int u,v;
        u = readIntSp(1,N);
        v = readIntLn(1,N);
        adj[u].emplace_back(v);
        adj[v].emplace_back(u);
    }
    dfs(1);
    for(int i = 1; i < 20; i++){
        for(int j = 1; j <= N; j++){
            dp[i][j]=dp[i-1][dp[i-1][j]];
        }
    }
    vector<bool>nodes(N+1);
    rep(i,1,500005){
        if(colnode[i].empty())continue;
        if(sz(colnode[i])==1){
            reqdnodes[i].emplace_back(colnode[i].front());
            continue;
        }
        colnode[i].reserve(2*sz(colnode[i]));
        for(auto j:colnode[i])nodes[j]=1;
        int k = sz(colnode[i]);
        rep(j,1,k){
            int tmp = lca(colnode[i][j-1],colnode[i][j]);
            if(!nodes[tmp]){
                nodes[tmp]=1;
                colnode[i].emplace_back(tmp);
            }
        }
        for(auto j:colnode[i])nodes[j]=0;
        sort(colnode[i].begin()+k,colnode[i].end(),comp);
        inplace_merge(colnode[i].begin(),colnode[i].begin()+k,colnode[i].end(),comp);
        stack<pii>stk;
        stk.emplace(0,INT_MAX);
        for(auto j:colnode[i]){
            while(outtime[j]>stk.top().S)stk.pop();
            temptree[stk.top().F].emplace_back(j,abs(height[j]-height[stk.top().F]));
            temptree[j].emplace_back(stk.top().F,abs(height[j]-height[stk.top().F]));
            stk.emplace(j,outtime[j]);
        }
        auto res = get_diameter(temptree[0][0].F,0);
        reqdnodes[i].emplace_back(res.best_res.S.F);
        reqdnodes[i].emplace_back(res.best_res.S.S);
        for(auto j:colnode[i]){
            temptree[j].clear();
        }
        temptree[0].clear();
        vi().swap(colnode[i]);
    }
    rep(i,1,M+1){
        int K,C;
        K = readIntSp(1,N);
        if(i==M) C = readInt(1,500000,EOF);
        else C = readIntLn(1,500000);
        int ans = -1;
        for(auto j:reqdnodes[C]){
            int LCA = lca(j,K);
            ans = max(ans,height[j]+height[K]-2*height[LCA]);
        }
        cout << ans << endl;
    }
}

int main(){
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    int t=1;
//  cin>>t;
    while(t--){
        solve();
    }
    return 0;
}
coding problems