### PROBLEM LINK:

**Author:** Sergey Nagin

**Tester:** Hiroto Sekido, Gerald Agapov

**Editorialist:** Ajay K. Verma

### DIFFICULTY:

Easy

### PREREQUISITES:

Depth first traversal

### PROBLEM:

Given a forest with N nodes, where each nodes is labeled with one of the two labels {0, 1}, find

- how many pair of label ‘1’ nodes exist which are connected by a path.
- how many nodes w are there in the forest which lie on the path between two label ‘1’ nodes.

### QUICK EXPLANATION:

Split the forest into trees and for each tree compute both results (1) and (2). The answer for the forest would be the sum of results of individual trees.

In a tree all nodes are connected with one another, hence (1) is equivalent to find how many label ‘1’ nodes exist in the tree. For (2) one needs to perform a depth-first traversal to compute the number of label ‘1’ nodes in the subtree rooted at a node x, for all x. Using these values one can calculate the number of nodes lying on the path between two label ‘1’ nodes (for more details see the next section);

### EXPLANATION:

The problem states that there can be at most one path between a pair of nodes, i.e., the graph has no cycles. This means the graph is a union of trees.

It is easy to see that the result for both (1) and (2) can be obtained by computing the result for individual tree and then adding them together. Hence, the problem reduces to solving the problem on a tree.

In a single tree all nodes are connected to one another, hence (1) is equivalent to computing the label ‘1’ nodes in the tree. If the number of such nodes is n, then the number of pairs will be (n * (n - 1))/2.

In order to compute the result for (2), we can iterate through all nodes w in the tree and check if it lies on the path between two label ‘1’ nodes. A naive way of considering all paths among label ‘1’ nodes will certainly run out of time as there can be O (N^{2}) such paths. Luckily, there is an easier way to check the same.

For this we make the tree rooted at some node r, and for each node x of the tree compute the number of label ‘1’ children in the subtree rooted at x. We denote this number by I(x). Now let us assume that the children nodes of w are y1, y2, …, yk. We can see that w will lie on the path between two label ‘1’ nodes if and only if one of the following conditions hold.

- 0 < I(w) < I®. In this case w lies on the path between a node in the subtree and a node outside the subtree both having label ‘1’. The first inequality ensures that there is a label ‘1’ node in the subtree rooted at w, and the second inequality ensures that there is at least one label ‘1’ node outside the subtree.
- label(w) = 1 and I(w) > 1. In this case subtree rooted at w has at least one node other than w with label ‘1’. The path between this node and w is the desired path.
- At least two of the I(y1), I(y2), …, I(yk) are non-zero. In this case the desired path will be between two nodes in the subtree rooted at w, which will pass through w.

This means that if we can compute I(x) for all nodes of the forest, then we can solve the problem easily. This can be achieved easily by traversing the nodes of the tree in reverse topological order, and using the following recurrence:

I(w) = label(w) + I(y1) + I(y2) + … + I(yk), where y1, y2, …, yk are children of w.

The following code partitions the forest into trees, and computes the I() of all nodes.

bool visited[N]; int I[N]; int parent[N]; int root_id[N]; // root node of the tree to which this node belongs to. int id = -1; // root of the current tree. // Perfroms the depth first traversal on the tree rooted at v, and computes // I() of all nodes in the tree. void do_dfs(int v) { visited[v] = true; I[v] = label[v]; root_id[v] = id; for (int x : neighbors[v]) { if (visited(x)) continue; parent[x] = v; do_dfs(x); I[v] += I[x]; } } // Partitions the forest into trees, and computes I() for al nodes. void partition_and_compute() { for (int i = 0; i < N; ++i) { if (!visited[i]) { id = i; do_dfs(i); } } }

### TIME COMPLEXITY:

O (N)

### AUTHOR’S AND TESTER’S SOLUTIONS:

Author’s solution will be put up soon.

Tester’s solution will be put up soon.