MEX on Trees Solution Codechef

by

MEX on Trees Solution Codechef

Consider a tree with N nodes, rooted at node 1. The value of the ith node (1≤i≤N) is denoted by Ai.

Chef defines the function MEX(u,v) as follows:
Let B denote the set of all the values of nodes that lie in the shortest path from u to v (including both u and v). Then, MEX(u,v) denotes the MEX of the set B. Find the maximum value of MEX(1,i), where 1≤i≤N.

Input Format

  • The first line of input contains a single integer T, denoting the number of test cases. The description of the T test cases follows.
  • The first line of each test case contains a single integer N – the number of nodes in the tree.
  • The second line of each test case contains N space separated integers A1,A2,…,AN. Ai represents the value of the ith node.
  • The following N−1 lines contain two space-separated integers X and Y, denoting that there exists an edge between the nodes X and Y.

Output Format

For each test case, output in a single line, the maximum value of MEX(1,i) where 1≤i≤N.

Constraints

  • 1≤T≤104
  • 1≤N≤105
  • 0≤Ai≤N
  • 1≤X,Y≤N
  • Sum of N over all test cases does not exceed 3⋅105.

Sample Input 1

3
3
0 1 2
1 2
1 3
4
0 1 2 3
1 2
2 3
1 4
5
1 2 3 4 0
1 2
2 3
2 4
1 5

Sample Output 1

2
3
2

Explanation

Test case 1:

MEX(1,1)=MEX({A1})=MEX({0})=1.
MEX(1,2)=MEX({A1,A2})=MEX({0,1})=2.
MEX(1,3)=MEX({A1,A3})=MEX({0,2})=1.
Thus, the answer is 2.

Test case 2:

MEX(1,1)=MEX({A1})=MEX({0})=1.
MEX(1,2)=MEX({A1,A2})=MEX({0,1})=2.
MEX(1,3)=MEX({A1,A2,A3})=MEX({0,1,2})=3.
MEX(1,4)=MEX({A1,A4})=MEX({0,3})=1.
Thus, the answer is 3.

Test case 3:

MEX(1,1)=MEX({A1})=MEX({1})=0.
MEX(1,2)=MEX({A1,A2})=MEX({1,2})=0.
MEX(1,3)=MEX({A1,A2,A3})=MEX({1,2,3})=0.
MEX(1,4)=MEX({A1,A2,A4})=MEX({1,2,4})=0.
MEX(1,5)=MEX({A1,A5})=MEX({1,0})=2.
Thus, the answer is 2.

SOLUTION

Program: MEX on Trees Solution Codechef in C++

#include <bits/stdc++.h>
using namespace std;
#define pii pair<int, int>
#define maxn 100010
int n, val[maxn], vis[maxn], ans[maxn], maxmex, par[maxn], vvis[maxn];
vector<int> adj[maxn];
set<pii > mex;
set<pii > :: iterator it;
void dfs(int pos) {
    vis[pos] = 1;
    bool flag = 1;
    if(mex.count(make_pair(1, val[pos]))) ans[pos] = ans[par[pos]];
    else {
        flag = 0;
        mex.erase(make_pair(0, val[pos]));
        mex.insert(make_pair(1, val[pos]));
        it = mex.begin();
        ans[pos] = it -> second;
    }
    maxmex = max(maxmex, ans[pos]);
    for(int i = 0; i < adj[pos].size(); i ++)
        if(!vis[adj[pos][i]]) par[adj[pos][i]] = pos, dfs(adj[pos][i]);
    if(!flag) {
        mex.erase(make_pair(1, val[pos]));
        mex.insert(make_pair(0, val[pos]));
    }
    return;
}
int main() {
    int tcase, st, en;
    scanf("%d", &tcase);
    while(tcase --) {
        scanf("%d", &n);
        memset(vis, 0, sizeof(vis));
        memset(ans, -1, sizeof(ans));
        mex.clear();
        maxmex = -1;
        for(int i = 1; i <= n; i ++) {
            scanf("%d", &val[i]);
            adj[i].clear();
        }
        for(int i = 0; i <= n; i ++) mex.insert(make_pair(0, i));
        for(int i = 1; i < n; i ++) {
            scanf("%d%d", &st, &en);
            adj[st].push_back(en);
            adj[en].push_back(st);
        }
        ans[1] = val[1] ? 0 : 1;
        mex.erase(make_pair(0, val[1]));
        mex.insert(make_pair(1, val[1]));
        par[1] = 1;
        dfs(1);
        printf("%d\n", maxmex);
    }
    return 0;
}

Program: MEX on Trees Solution Codechef in C++

#include<bits/stdc++.h>
using namespace std;
#define Mx 100100
bool vis [Mx];
vector<int> v[Mx];
set<int> k;
int cnt[Mx], a[Mx];
int tcase, n, x, y, maxe;
void dfs(int x) {
    vis[x] = 1;
    cnt[a[x]]++;
    k.erase(a[x]);
    maxe = max(maxe, *(k.begin()));
    for (auto s : v[x]) {
        if (vis[s]) continue;
        dfs(s);
    }
    cnt[a[x]]--;
    if (!cnt[a[x]]) k.insert(a[x]);
    return ;
}
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
    for (cin >> tcase; tcase--;) {
        cin >> n;
        k.clear();
        k.insert(0);
        for (int i = 1; i <= n; i++) {
            cin >> a[i];
            vis[i] = cnt[i] = 0;
            v[i].clear();
            k.insert(i);
        }
        for (int i = 1; i < n; i++) {
            cin >> x >> y;
            v[x].push_back(y);
            v[y].push_back(x);
        }
        maxe = 0;
        dfs(1);
        cout << maxe << endl;
    }
    return 0;
}

Codechef March Long Challenge 2022 Solutions

Leave a Comment

13 − 2 =