HackerEarth Little Shino and K Ancestor problem solution

In this HackerEarth Little Shino and K Ancestor problem solution Assume that you are given an undirected rooted tree with N nodes and an integer K. Node 1 is the root of the tree. Each node is uniquely numbered from 1 to N. Additionally, each node also has a color and the color is an integer value.

Note: Different nodes can have the same color.

For each node, you are required to find the Kth closest ancestor from that node which has the same color.

HackerEarth Little Shino and K Ancestor problem solution.

#include <bits/stdc++.h>

using namespace std;

const int MAX = 1e6 + 5;
int a[MAX], k, n, ans[MAX];
vector <int> v[MAX];
vector <int> adj[MAX];

void dfs(int from, int par) {
    if (v[a[from]].size() >= k) {
        int l = v[a[from]].size();
        ans[from] = v[a[from]][l - k];
    }
    else {
        ans[from] = -1;
    }
    for (int to : adj[from]) {
        if (to != par) {
            v[a[from]].push_back(from);
            dfs(to, from);
            v[a[from]].pop_back();
        }
    }
}

int main(int argc, char* argv[])
{
    if(argc == 2 or argc == 3) freopen(argv[1], "r", stdin);
    if(argc == 3) freopen(argv[2], "w", stdout);
    ios::sync_with_stdio(false);
    int x, y;
    cin >> n >> k;
    for (int i = 1; i <= n; i++) {
        cin >> a[i];
    }
    for (int i = 0; i < n-1; i++) {
        cin >> x >> y;
        adj[x].push_back(y);
        adj[y].push_back(x);
    }
    dfs(1, -1);
    for (int i = 1; i <= n; i++) {
        if (i < n) cout << ans[i] << ' ';
        else cout << ans[i] << endl;
    }
    return 0;
}

Second solution

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#include <algorithm>
#include <set>
#include <map>
#include <cmath>
#include <ctime>
#include <functional>
#include <sstream>
#include <fstream>
#include <valarray>
#include <complex>
#include <queue>
#include <cassert>
#include <bitset>
using namespace std;

#ifdef LOCAL
#define debug_flag 1
#else
#define debug_flag 0
#endif

  template <class T1, class T2 >
std::ostream& operator << (std::ostream& os, const pair<T1, T2> &p) 
{
  os << "[" << p.first << ", " << p.second << "]";
  return os;
}

  template <class T >
std::ostream& operator << (std::ostream& os, const std::vector<T>& v) 
{
  os << "[";
  bool first = true;
  for (typename std::vector<T>::const_iterator it = v.begin(); it != v.end(); ++it)
  {
    if (!first)
      os << ", ";
    first = false;
    os << *it;
  }
  os << "]";
  return os;
}

#define dbg(args...) { if (debug_flag) { _print(_split(#args, ',').begin(), args); cerr << endl; } else { void(0);} }

vector<string> _split(const string& s, char c) {
  vector<string> v;
  stringstream ss(s);
  string x;
  while (getline(ss, x, c))
    v.emplace_back(x);
  return v;
}

void _print(vector<string>::iterator) {}
template<typename T, typename... Args>
void _print(vector<string>::iterator it, T a, Args... args) {
  string name = it -> substr((*it)[0] == ' ', it -> length());
  if (isalpha(name[0]))
    cerr << name  << " = " << a << " ";
  else
    cerr << name << " ";
  _print(++it, args...);
}

typedef long long int int64;

const int N = (int)1e6 + 100;

int n, k;
int c_list[N];
vector<int> s_list[N];
int ans_list[N];
vector<int> graph[N];

void dfs(int v, int p)
{
  int color = c_list[v];

  s_list[color].push_back(v);

  int s_size = (int)s_list[color].size();
  if (s_size > k) {
    ans_list[v] = s_list[color][s_size - k - 1] + 1;
  } else {
    ans_list[v] = -1;
  }

  for (int to : graph[v]) {
    if (to == p) {
      continue;
    }
    dfs(to, v);
  }

  s_list[color].pop_back();
}

int main()
{
#ifdef LOCAL
  freopen ("input.txt", "r", stdin);
#endif

  scanf("%d%d", &n, &k);
  for (int i = 0; i < n; i++) {
    scanf("%d", &c_list[i]);
  }
  for (int i = 0; i < n - 1; i++) {
    int a, b;
    scanf("%d%d", &a, &b);
    a--;
    b--;
    graph[a].push_back(b);
    graph[b].push_back(a);
  }

  dfs(0, -1);

  for (int i = 0; i < n; i++) {
    printf("%d ", ans_list[i]);
  }
  printf("\n");

  return 0;
}