HackerEarth Flight Tickets problem solution

In this HackerEarth Flight Tickets problem solution There are N cities in the country Byteland. Each city has some special value associated with it. All the cities are connected to one another either via direct flights or a series of connecting flights. There are a total of N – 1 flights available and it is always possible to reach any city from any other city thus the structure of flights between the cities in Byteland forms a tree. The fare of flight between two cities is given by the distance between the two cities in the tree.

Now Rohan books flights between the two cities if and only if their special value is equal. He wants to book the cheapest and distinct K flights. He follows the following rules while booking flights-

• Two flights are distinct if their start city or destination city is different from other booked flights.
• Note that since Byteland is mysterious, booking of flight from the same city to the same city is also possible with a cost of 0.
• If Rohan books a flight from the city A to city B then he can’t book any flight from the city B to city A.
• In the output, you need to print the price of the last flight that he would book. If it is impossible to book K flights then you need to print -1.

#include<bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define LL long long int
#define M 1000000007
#define endl "n"
#define eps 0.00000001
LL pow(LL a,LL b,LL m){LL x=1,y=a;while(b > 0){if(b%2 == 1){x=(x*y);if(x>m) x%=m;}y = (y*y);if(y>m) y%=m;b /= 2;}return x%m;}
LL gcd(LL a,LL b){if(b==0) return a; else return gcd(b,a%b);}
LL gen(LL start,LL end){LL diff = end-start;LL temp = rand()%start;return temp+diff;}
using namespace __gnu_pbds;
using namespace std;
typedef tree<
  pair<int,int>,
  null_type,
  less<pair<int,int>>,
  rb_tree_tag,
  tree_order_statistics_node_update>
ordered_set;
/********************** Centroid Decomposition Code **************************/
#define maxn 30001
int c[maxn + 1];
int parent[maxn + 1];//this stores the structure of the centroid tree
set<int> graph[maxn + 1];//original graph
int tree_sz = 0;
int sz[maxn + 1];
int find_centroid(int node,int p = -1){
    for(int i: graph[node]){
        if(i != p && sz[i] > tree_sz / 2){
            return find_centroid(i , node);
        }
    }
    return node;
}

void update_subtrees(int node,int p = -1){
    ++tree_sz;
    sz[node] = 1;
    for(int i: graph[node]){
        if(i != p){
            update_subtrees(i , node);
            sz[node] += sz[i];
        }
    }
}

void centroid_decomposition(int node,int p = -1){
    tree_sz = 0;
    update_subtrees(node);
    int centroid = find_centroid(node);
    if(p == -1){
        p = centroid;
        parent[centroid] = centroid;
    }
    else{
        parent[centroid] = p;
    }
    for(int i: graph[centroid]){
        graph[i].erase(centroid);
        centroid_decomposition(i , centroid);
    }
    graph[centroid].clear();
}
/******************************************************************************/
int LCA[20][maxn + 1];
int dist[maxn + 1];
void dfs_distance(int node,int p = -1){
    for(int i: graph[node]){
        if(i != p){
            LCA[0][i] = node;
            dist[i] = dist[node] + 1;
            dfs_distance(i , node);
        }
    }
}


int lca(int u,int v){
    if(dist[u] < dist[v])
        swap(u , v);
    int len = dist[u] - dist[v];
    for(int i = 15;i >= 0; i--){
        if(len & (1<<i)){
            u = LCA[i][u];
        }
    }
    if(u == v)
        return u;
    for(int i = 15; i >= 0; i--){
        if(LCA[i][u] != LCA[i][v]){
            u = LCA[i][u];
            v = LCA[i][v];
        }
    }
    return LCA[0][u];
}

int distance(int u,int v){
    return dist[u] + dist[v] - 2 * dist[lca(u, v)];
}

void preprocess_lca(int n){
    for(int i = 1; i <= 15; i++){
        for(int j = 1; j <= n; j++){
            LCA[i][j] = LCA[i - 1][LCA[i - 1][j]];
        }
    }
}

void preprocess_distance(int n){
    dfs_distance(1);
    preprocess_lca(n);
}

int pekka = 0;
int val[maxn + 1];
vector<int> g[maxn + 1];
ordered_set m[maxn + 1];
int D[100001];
int cur_root;
void add_subtree(int node){
    stack<int> s;
    s.push(node);
    while(s.size()){
        int node = s.top();
        s.pop();
        if(node != cur_root){
            m[c[node]].insert({D[node] = distance(cur_root , node) , ++pekka});
            val[node] = pekka;
        }
        for(int i: g[node]){
            if(i != parent[node]){
                s.push(i);
            }
        }
    }
}
void remove_subtree(int node){
    stack<int> s;
    s.push(node);
    while(s.size()){
        int node = s.top();
        s.pop();
        if(node != cur_root)
            m[c[node]].erase({D[node] , val[node]});
        for(int i: g[node]){
            if(i != parent[node]){
                s.push(i);
            }
        }
    }
}
int total = 0;
int get_ans(int node,int sz){
    int temp = 0;
    int rem = distance(cur_root , node);
    if(rem <= sz && c[cur_root] == c[node]){
        ++total;
    }
    rem = sz - rem;
    if(rem >= 0 && m[c[node]].size() > 0)
        temp = temp + max(0 , (int)m[c[node]].order_of_key({rem + 1 , -1}));
    for(int i: g[node]){
        if(i != parent[node]){
            temp = temp + get_ans(i , sz);
        }
    }
    return temp;
}

int check(int sz,int n){
    total = 0;
    int extra = 0;
    for(int node = 1; node <= n; node++){
        cur_root = node;
        add_subtree(node);
        for(int i: g[node]){
            if(i == parent[node])
                continue;
            remove_subtree(i);
            extra = extra + get_ans(i , sz);
        }
        remove_subtree(node);
    }   
    total += extra + n;
    return total;
}


int main()
    {
       // ios_base::sync_with_stdio(0);
        int n , k;
        scanf("%d%d" , &n , &k);
        for(int i = 1; i < n; i++){
            int u , v;
            scanf("%d%d" , &u , &v);
            graph[u].insert(v);
            graph[v].insert(u);
        }
        set<int> hash_;
        map<int,int> hash__;
        int hash_ctr = 0;
        for(int i = 1; i <= n; i++){
            scanf("%d" , &c[i]);
            hash_.insert(c[i]);
        }
        for(int i: hash_){
            hash__[i] = ++hash_ctr;
        }
        for(int i = 1; i <= n; i++){
            c[i] = hash__[c[i]];
        }
        preprocess_distance(n);
        centroid_decomposition(1);
        int root;   
        for(int i = 1; i <= n; i++){
            if(parent[i] != i){
                g[i].push_back(parent[i]);
                g[parent[i]].push_back(i);
            }
            else{
                root = i;
            }
        }
        int l = 0 , r = n;
        int ans = n + 1;
        while(l <= r){
            int mid = (l + r) / 2;
            if(check(mid , n) >= k){
                ans = mid;
                r = mid - 1;
            }
            else{
                l = mid + 1;
            }
        }
        if(ans == n + 1){
            printf("-1");
        }
        else{
            printf("%d" , ans);
        }
    }