AtCoder Educational DP - V Subtrees

The link to problem is here: https://atcoder.jp/contests/dp/tasks/dp_v
I followed the same idea as the one mentioned in the editorial, but after thorough debugging, I still cannot figure out why my code gets WA:

import java.io.*;
import java.util.*;

public class VSubtree {
    static long[][] suffix; 
    static long[] dp;
    static ArrayList<Integer>[] adj;
    static int mod;
    
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        mod = Integer.parseInt(st.nextToken());
        
        suffix = new long[N][];
        
        adj = new ArrayList[N];
        for(int i = 0; i<N; i++) adj[i] = new ArrayList();
        for(int i = 0; i<N-1; i++){
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken())-1;
            int v = Integer.parseInt(st.nextToken())-1;
            adj[u].add(v);
            adj[v].add(u);
        }
        dp = new long[N];
        dfs(0,-1);
        dfs1(0,-1,0);
        
        for(long i : dp){
            System.out.println(i);
        }
    }
    
    static long dfs(int curr, int last){        
        suffix[curr] = new long[adj[curr].size()+1];
        long right = 1;
        suffix[curr][adj[curr].size()] = 1; //idx out of bounds
        
        for(int i = adj[curr].size()-1; i>=0; i--){
            int n = adj[curr].get(i);
            if(n == last) continue;
            right *= (dfs(n, curr)+1);
            right %= mod;
            suffix[curr][i] = right;
        }
        dp[curr] = right;
        return dp[curr];
    }
    
    static void dfs1(int curr, int last, long above){
        long left = above+1;
        
        dp[curr] *= left;
        dp[curr] %= mod;
        
        for(int i = 0; i<adj[curr].size(); i++){
            int n = adj[curr].get(i);
            if(n == last) continue;
            
            long temp = dp[n]+1;
            
            dfs1(n, curr, (left*suffix[curr][i+1])%mod);
                        
            left *= temp;
            left %= mod;
        }
    }
}

Any help would be appreciated!

Your code should be formatted like so:

import java.io.*;
import java.util.*;

public class VSubtree {
       static long[][] suffix; 
        static long[] dp;
       static ArrayList<Integer>[] adj;
       static int mod;
    
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        mod = Integer.parseInt(st.nextToken());
        
        suffix = new long[N][];
        
        adj = new ArrayList[N];
        for(int i = 0; i<N; i++) adj[i] = new ArrayList();
        for(int i = 0; i<N-1; i++){
            st = new StringTokenizer(br.readLine());
            int u = Integer.parseInt(st.nextToken())-1;
            int v = Integer.parseInt(st.nextToken())-1;
            adj[u].add(v);
            adj[v].add(u);
        }
        dp = new long[N];
        dfs(0,-1);
        dfs1(0,-1,0);
        
        for(long i : dp){
            System.out.println(i);
        }
    }
    
    static long dfs(int curr, int last){        
        suffix[curr] = new long[adj[curr].size()+1];
        suffix[curr][adj[curr].size()] = 1; //idx out of bounds
        
        for(int i = adj[curr].size()-1; i>=0; i--){
            int n = adj[curr].get(i);
            if(n == last) continue;
            suffix[curr][i] = (dfs(n, curr)+1)*(suffix[curr][i+1])%mod;
        }
        dp[curr] = suffix[curr][0];
        return dp[curr];
    }
    
    static void dfs1(int curr, int last, long above){
        long left = above+1;
        
        dp[curr] *= left;
        dp[curr] %= mod;
        
        for(int i = 0; i<adj[curr].size(); i++){
            int n = adj[curr].get(i);
            if(n == last) continue;
            
            long temp = dp[n]+1;
            
            dfs1(n, curr, (left*suffix[curr][i+1])%mod);
                        
            left *= temp;
            left %= mod;
        }
    }
}

but after thorough debugging, I still cannot figure out why my code gets WA:

If you have thoroughly debugged, how are you seem to be failing the first sample case? :confused:

Thanks for pointing that out.

I modified my code a bit before posting it here(to make it more readable), and I didn’t think it would make a difference. Here was my original submission:

Considering that you’re failing the last sample case, I’ll leave it up to you to debug it … Try

  • Finding a smaller test case for which your code gives different output than an AC solution (change the mod if that helps)
  • Simulating your code by hand on that test case (add print statements as necessary).

Refer to the debugging module linked here for more info.

1 Like

I thought I had passed the last test case. Thank you for letting me know. I was able to find my mistake(leaving suffix value as 0 when node is parent) after having a sample case.

Benq Sir please look upon my code, It passed only 9 TC, rest are WA

/ * Author :- Himanshu Shekhar , IIIT Bhagalpur */

#include <bits/stdc++.h>
using namespace std;
const int MOD = 1e9 + 7;
const long long INF = 1e15;
const int MAX_N = 200 * 1000 + 13;

# define ll long long

vector<vector<int>> graph;

ll a,b,c,d,e,f,g,h,i,j,k,l,m,n,o,p,q,r,s,t,u,v,w,x,y,z,sum,maxi;
bool flag;
string str;


// IN-OUT DP or UP-DOWN DP


ll in[MAX_N]; // // number of ways to colour the nodes outside the subtree of i
ll out[MAX_N]; // number of ways to colour the nodes outside the subtree of i
void dfs(int source, int parent){
  
  in[source]=1;  // choice 1: for painting black

  for(auto it: graph[source]){
    if(it!=parent){
      dfs(it,source);
      in[source]=(in[source]*in[it]);
    }
  }
  in[source]++; // for painting all whites
}


void dfs2(int source, int parent){

      vector<ll> pre,suff;

      for(auto it: graph[source]){
        if(it!=parent)
        pre.push_back(in[it]);
        suff.push_back(in[it]);
      }

      int sz = pre.size();

      for(int i = 1; i<sz; i++){
        pre[i]=(pre[i]*pre[i-1]);
      }

      for(int i = sz-2; i>=0; i--){
        suff[i]=(suff[i]*suff[i+1]);
      }

      int cnt = 0;


  for(auto it: graph[source]){
      if(it==parent) continue;

      ll t2 = in[it];

          // re-rooting from node to ch
    ll l=1, r=1;    
    if(i-1>=0)
      l=pre[i-1];
    if(i+1<sz)
     r=suff[i+1];

    out[it]=(l*r*out[source]); 
    out[it]++; 
    dfs2(it,source);
    cnt++;     
  }
}


void solve(){
    
    cin>> n >> m;
    graph.resize(n+1);

    // constructing the graph
    for(int i = 1; i<n; i++){
            cin >> a >> b;
            a--;
            b--;
            graph[a].push_back(b);
            graph[b].push_back(a);
    }

    fill(in,in+MAX_N,1);
    fill(out,out+MAX_N,1);

    dfs(0,-1); // calculation in[i] for each node i
    
    dfs2(0,-1); // calculation out[i] for each node i

    for(int i = 0; i<n; i++){
      cout << ((in[i]-1)*out[i])%m << endl; 
    }

    cout << endl;   

       fflush(stdout);
       cout << flush;
}

int main() {

    #ifndef ONLINE_JUDGE
//for getting input from input.txt
freopen("input1.txt", "r", stdin);
//for writing output to output.txt
freopen("output.txt", "w", stdout);
#endif

ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(NULL);

   solve();
       return 0;
}