D. 0-1-Tree

Problem based on DP on trees.

Java/D. 0-1-Tree

@@ -0,0 +1,184 @@
+Problem : 
+
+You are given a tree (an undirected connected acyclic graph) consisting of n vertices and n−1 edges. 
+A number is written on each edge, each number is either 0 (let's call such edges 0-edges) or 1 (those are 1-edges).
+
+Let's call an ordered pair of vertices (x,y) (x≠y) valid if, while traversing the simple path from x to y, 
+we never go through a 0-edge after going through a 1-edge. 
+
+Your task is to calculate the number of valid pairs in the tree.
+
+Input
+The first line contains one integer n (2≤n≤200000) — the number of vertices in the tree.
+
+Then n−1 lines follow, each denoting an edge of the tree. Each edge is represented by three integers xi, yi and ci (1≤xi,yi≤n, 0≤ci≤1, xi≠yi) — the vertices connected by this edge and the number written on it, respectively.
+
+It is guaranteed that the given edges form a tree.
+
+Output
+Print one integer — the number of valid pairs of vertices.
+
+
+Topics : DFS, DP on Trees
+
+Solution in JAVA:
+
+
+import java.util.*;
+import java.lang.*;
+import java.io.*;
+
+
+public class Solution {
+
+    static class Edge {
+        int v;
+        int c;
+
+        Edge(int v, int c) {
+            this.v = v;
+            this.c = c;
+        }
+    }
+
+    static int MAX = 200005;
+    static List<Edge>[] graph = new List[MAX];
+
+    // dp[i][0] is the number of paths having only 0-edges or {0000...11111}.
+    // dp[i][1] is the number of paths having only 1-edges.
+    static int[][] dp = new int[MAX][2];
+
+    public static void main(String[] args) throws java.lang.Exception {
+        out = new PrintWriter(new BufferedOutputStream(System.out));
+        sc = new FastReader();
+
+        int test = 1;
+        for (int t = 1; t <= test; t++) {
+            solve();
+        }
+        out.close();
+    }
+
+    private static void solve() {
+        int n = sc.nextInt();
+        for (int i = 1; i <= n; i++) {
+            graph[i] = new ArrayList<>();
+        }
+
+        for (int i = 1; i < n; i++) {
+            int u = sc.nextInt();
+            int v = sc.nextInt();
+            int color = sc.nextInt();
+            graph[u].add(new Edge(v, color));
+            graph[v].add(new Edge(u, color));
+        }
+
+        dfs(1, 0);
+
+        // re-rooting the root since our path can start from any node.
+        dfs2(1, 0);
+
+        long pairs = 0;
+        for (int i = 1; i <= n; i++) {
+            pairs += dp[i][0] + dp[i][1];
+        }
+
+        out.println(pairs);
+    }
+
+    private static void dfs2(int currNode, int parent) {
+        for (Edge adjacentEdge : graph[currNode]) {
+            int to = adjacentEdge.v;
+            if (to == parent) {
+                continue;
+            }
+            if (adjacentEdge.c == 0) {
+                dp[to][0] = dp[currNode][0] + dp[currNode][1] - dp[to][1];
+            }else {
+                dp[to][1] = dp[currNode][1];
+            }
+            dfs2(to, currNode);
+        }
+    }
+
+    private static void dfs(int currNode, int parent) {
+        for (Edge adjacentEdge : graph[currNode]) {
+            int to = adjacentEdge.v;
+            if (to == parent) {
+                continue;
+            }
+            dfs(to, currNode);
+            if (adjacentEdge.c == 0) {
+                // path can be just 0
+                // path can be 00000...
+                // path can be 0000...11111
+                dp[currNode][0] += dp[to][0] + dp[to][1] + 1;
+            }else {
+                // path can be just 1.
+                // path can be 11111...
+                dp[currNode][1] += dp[to][1] + 1;
+            }
+        }
+    }
+
+
+    public static FastReader sc;
+    public static PrintWriter out;
+    static class FastReader
+    {
+        BufferedReader br;
+        StringTokenizer str;
+
+        public FastReader()
+        {
+            br = new BufferedReader(new
+                    InputStreamReader(System.in));
+        }
+
+        String next()
+        {
+            while (str == null || !str.hasMoreElements())
+            {
+                try
+                {
+                    str = new StringTokenizer(br.readLine());
+                }
+                catch (IOException  lastMonthOfVacation)
+                {
+                    lastMonthOfVacation.printStackTrace();
+                }
+            }
+            return str.nextToken();
+        }
+
+        int nextInt()
+        {
+            return Integer.parseInt(next());
+        }
+
+        long nextLong()
+        {
+            return Long.parseLong(next());
+        }
+
+        double nextDouble()
+        {
+            return Double.parseDouble(next());
+        }
+
+        String nextLine()
+        {
+            String str = "";
+            try
+            {
+                str = br.readLine();
+            }
+            catch (IOException lastMonthOfVacation)
+            {
+                lastMonthOfVacation.printStackTrace();
+            }
+            return str;
+        }
+    }
+
+}