Tree Permutations - CodeChef April Long Challenge Solution | AskTheCode

Updated: Apr 12

CodeChef April Long Challenge Solution | Tree Permutations (TREEPERM) solution in Java | AskTheCode

Problem:

You are given a tree with N nodes (numbered 1 through N), rooted at node 1. For each valid i, node i has a value ai written on it.


An undirected simple path between any two nodes u and v is said to be vertical if u=v or u is an ancestor of v or v is an ancestor of u. Let's define a vertical partition of the tree as a set of vertical paths such that each node belongs to exactly one of these paths.


You are also given a sequence of N integers b1,b2,…,bN. A vertical partition is good if, for each of its paths, we can permute the values written on the nodes in this path, and this can be done in such a way that we reach a state where for each valid i, the value written on node i is bi.


The difficulty of your task is described by a parameter S. If S=1, your task is only to determine whether at least one good vertical partition exists. If S=2, you are required to find the number of good vertical partitions modulo 1,000,000,007 (10^9+7).

Input:

  • The first line of the input contains a single integer T denoting the number of test cases. The description of T test cases follows.

  • The first line of each test case contains two space-separated integers N and S.

  • Each of the next N−1 lines contains two space-separated integers u and v denoting that nodes u and v are connected by an edge.

  • The next line contains N space-separated integers a1,a2,…,aN.

  • The line after that contains N space-separated integers b1,b2,…,bN.

Output:

For each test case, print a single line containing one integer:

  • If S=1, this integer should be 1 if a good vertical partition exists or 0 if it does not exist.

  • If S=2, this integer should be the number of good vertical partitions modulo 1,000,000,007 (10^9+7).

Example Input:

4 3 2 1 2 2 3 4 5 6 4 6 5 6 2 1 2 1 3 2 4 3 5 3 6 10 20 30 40 50 60 10 40 50 20 30 60 6 1 1 2 1 3 2 4 3 5 3 6 10 20 30 40 50 60 10 40 50 20 30 60 1 2 1 2

Example Output:

2 3 1 0

Code:

In Java

import java.io.*;
import java.util.*;
class TestClass {
    static int mod = 1000000007;
        static ArrayList<Integer> al[];
        static int t_a[]=new int[1000005],t_b[]=new int[1000005],con[]=new int[1000005],a[]=new int[1000005],b[]=new int[1000005],par[],h[];
        static boolean vis[];
         static int determiner;
       static ArrayList<Integer> sset;
       static PriorityQueue<pair> pq;
       static void make_it(int nde,int d)
        {
            vis[nde]=true;
            h[nde]=d;

            boolean is_it=true;

            for(int i=0;i<al[nde].size();i++)
            {
                if(!vis[al[nde].get(i)])
                {
                    par[al[nde].get(i)]=nde;
                    make_it(al[nde].get(i),d+1);
                    is_it=false;
                }
            }

            if(is_it)
                pq.add(new pair(d,nde));
        }

       static void clean(int nde)
        {
            con[a[nde]]=0;
            con[b[nde]]=0;
            t_a[a[nde]]=0;
            t_a[b[nde]]=0;
            t_b[a[nde]]=0;
            t_b[b[nde]]=0;
        }

       static int make_set(int nde)
        {
            t_b[b[nde]]++;
            t_a[a[nde]]++;

            if(t_a[a[nde]]==t_b[a[nde]] && con[a[nde]]!=0)
            {
                con[a[nde]]--;
                determiner--;
            }
            else if(con[a[nde]]==0)
            {
                con[a[nde]]++;
                determiner++;
            }
            if(t_a[b[nde]]==t_b[b[nde]] && con[b[nde]]!=0)
            {
                con[b[nde]]--;
                determiner--;
            }
            else if(con[b[nde]]==0)
            {
                con[b[nde]]++;
                determiner++;
            }
            vis[nde]=true;
            sset.add(nde);

            if(determiner==0)
            {
                if(vis[par[nde]]==false && nde!=1)
                    pq.add(new pair(h[par[nde]],par[nde]));

                clean(nde);

                return 1;
            }

            if(nde==1)
            {
                clean(nde);
                return 0;
            }

            if(vis[par[nde]]==false)
            {
                if(make_set(par[nde])==1)
                {
                    clean(nde);
                    return 1;
                }
            }
            clean(nde);

            return 0;
        }
     
    public static void main(String args[] ) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int t = Integer.parseInt(br.readLine());
        while(t-->0)
        {
            String ss[] = br.readLine().split(" ");
            int n = Integer.parseInt(ss[0]);
            int s = Integer.parseInt(ss[1]);
            al = new ArrayList[n+1];
            pq = new PriorityQueue<>(new Comparator<pair>() {
                @Override
                public int compare(pair o1, pair o2) {
                    if(o1.x==o2.x)
                        return (int)(o2.y-o1.y);
                    return (int)(o2.x-o1.x);
                }
            });
            vis = new boolean[n+1];
            par = new int[n+1];
            h = new int[n+1];
            sset= new ArrayList<>();
            for(int i=0;i<=n;i++)
                al[i]=new ArrayList<>();

            for(int i=0;i<n-1;i++)
            {
                String ed[] = br.readLine().split(" ");
                int x = Integer.parseInt(ed[0]);
                int y = Integer.parseInt(ed[1]);
                al[x].add(y);
                al[y].add(x);
            }
            String aa[] = br.readLine().split(" ");
            for(int i=1;i<=n;i++)
            {
                a[i] = Integer.parseInt(aa[i-1]);
            }
            String bb[] = br.readLine().split(" ");
            for(int i=1;i<=n;i++)
            {
                b[i] = Integer.parseInt(bb[i-1]);
            }



            make_it(1,1);
            Arrays.fill(vis,false);


            boolean correct=true;

            ArrayList<ArrayList<Integer>> sets = new ArrayList<>();

            while(!pq.isEmpty())
            {
                pair leaf=pq.poll();


                if(vis[(int)leaf.y]==false)
                {
                    determiner=0;
                    sset.clear();

                    if(make_set((int)leaf.y)==0)
                    {
                        correct=false;
                        break;
                    }
                    else
                    {
                        sets.add(new ArrayList<>());
                        sets.get(sets.size()-1).addAll(sset);

                    }
                }
            }
            if(correct==false)
            {
                System.out.println(0);
                continue;
            }

            

            long a=1;
            int x=sets.size();

            for(int i=0;i<x;i++)
            {
                int u=sets.get(i).get(0);
                int l=u;

                int sze=sets.get(i).size();

                for(int j=1;j<sze;j++)
                {
                    if(h[sets.get(i).get(j)]>h[l])
                        l=sets.get(i).get(j);
                    if(h[sets.get(i).get(j)]<h[u])
                        u=sets.get(i).get(j);
                }

                int cnt=0;
                for(Integer c:al[l])
                {
                    if(c!=par[l])
                        cnt++;
                }

                a=(a*(cnt+1))%mod;
            }
            if(s==1)
            {
                System.out.println(1);
                continue;
            }
            System.out.println(a);

        }
    
    }
    static class pair{
        long x;
        long y;
        pair(long x,long y)
        {
            this.x=x;
            this.y=y;
        }
    }
}

In C++

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll MOD = 1000000007;
vector <ll> Graph[100001] ;
ll A[1000001];
ll B[1000001];
ll con[1000001];
ll a[100001];
ll b[100001];
ll par[100001];
ll h[100001];
ll vis[100001];
vector<ll> ssat;
priority_queue<pair<ll,ll>> pq;
ll n,s;
ll TEMP;
 
void FIRST(){
    for(ll i=0;i<n+1;i++){
        Graph[i].clear();
        a[i]=0;
        b[i]=0;
        par[i]=0;
        h[i]=0;
    }
    while(!pq.empty()){
        pq.pop();
    }
}
void CLEAR(ll sup){
    con[a[sup]]=0;
    con[b[sup]]=0;
    A[a[sup]]=0;
    B[a[sup]]=0;
    A[b[sup]]=0;
    B[b[sup]]=0;
}
void solve1(ll SUP,ll d){
    vis[SUP]++;
    h[SUP]=d;
    bool is_it=true;
    for(ll chd:Graph[SUP]){
        if(!vis[chd]){
            par[chd]=SUP;
            solve1(chd,d+1);
            is_it=false;
        }
    }
    if(is_it==true){
        pq.push({d,SUP});
    }
}
int solve2(ll sup){
    B[b[sup]]++;
    A[a[sup]]++;
    if( A[a[sup]]== B[a[sup]] && con[a[sup]]!=0){
        con[a[sup]]--;
        TEMP--;
    }else if(con[a[sup]]==0){
        con[a[sup]]++;
        TEMP++;
    }
    if(A[b[sup]]==B[b[sup]] && con[b[sup]]!=0){
        con[b[sup]]--;
        TEMP--;
    }else if(con[b[sup]]==0){
        con[b[sup]]++;
        TEMP++;
    }
    vis[sup]++;
    ssat.push_back(sup);
    if(TEMP==0){
        if(vis[par[sup]]==0 && sup!=1)
            pq.push(make_pair(h[par[sup]],par[sup]));
     CLEAR(sup);
        return 1;
    } 
    if(sup==1){
          //cout<<"correct temp"<<endl;
         CLEAR(sup);
            return 0;
    }
    if(vis[par[sup]]==0){
        if(solve2(par[sup])==1){
         CLEAR(sup);
            return 1;
        }
    }
 CLEAR(sup);
    return 0;
}
int sol(){
    cin>>n>>s;
    for(ll i=0;i<n-1;i++){
        ll u,v;
        cin>>u>>v;
        Graph[u].push_back(v);
        Graph[v].push_back(u);
    }
    for(ll i=1;i<n+1;i++){
        cin>>a[i];
    }
      for(ll i=1;i<n+1;i++){
        cin>>b[i];
    }
      for(ll i=1;i<n+1;i++){
        vis[i]=0;
        par[i]=0;
        h[i]=0;
    }
    solve1(1,1);
    for(ll i=1;i<n+1;i++)
       vis[i]=0;
    bool correct=true;
    vector<vector<ll>> sets;
    while(!pq.empty()){
        pair<ll,ll> leaf=pq.top();
        pq.pop();
        if(vis[leaf.second]==0){
            TEMP=0;
            ssat.clear();
            int tem=solve2(leaf.second);
           // cout<<correct<<"correct"<<endl;
            if(tem==0){
               // cout<<"testing"<<endl;
                correct=false;
                break;
            }else
                sets.push_back(ssat);          
        }
    }
    
    if(correct==false){
        cout<<0<<endl;
        return 0;
    }
    if(s==1){
        cout<<1<<endl;
        return 0;
    }
    ll a=1;
    ll x=sets.size();
    for(ll i=0;i<x;i++){
        ll u=sets[i][0];
        ll l=sets[i][0];
        ll sze=sets[i].size();
        for(ll j=1;j<sze;j++){
            if(h[sets[i][j]]>h[l]){
                l=sets[i][j];
            }
            if(h[sets[i][j]]<h[u]){
                u=sets[i][j];
            }
        }
        ll cnt=0;
        for(ll c:Graph[l]){
            if(c!=par[l]){
                cnt++;
            }
        }
        a=(a*(cnt+1))%MOD;
    }
    cout<<a<<endl;  
    return 0;
}
int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
  ll t ;
  cin>>t;
  while(t--){
      sol();
      FIRST();
  }
  return 0;
}
839 views4 comments

Recent Posts

See All

Golf CodeChef Solution in Java - AskTheCode

CodeChef May Long Challenge Solution | Golf (LKDNGOLF) solution | AskTheCode Golf (LKDNGOLF) solution Problem: It's a lockdown. You’re bored in your house and are playing golf in the hallway. The hall

Solubility - CodeChef Solution in Java and C++| AskTheCode

CodeChef May Long Challenge Solution | Solubility (SOLBLTY) solution | AskTheCode Solubility (SOLBLTY) solution Problem: Suppose for a unit rise in temperature, the solubility of sugar in water increa

Valid Paths CodeChef Solution in C++ | AskTheCode

CodeChef May Long Challenge Solution | Valid Paths (VPATH) solution | AskTheCode Valid Paths (VPATH) CodeChef solution Problem: You are given a tree with N nodes numbered from 1 to N. A set S of nodes