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.