Last Updated: 15 Oct, 2021

Subtree Count

Moderate

Asked in company

Solve now

Problem statement

You are given an array ‘arr’ of the size ‘N’. The array is formed using a tree with the generator function. You have to obtain the tree which was used to generate the array and return the number of nodes in subtree of each node.

Below is the code given used to create the array - :

Function generate(node):

    arr.Add(node.id)
    visited[node.id] = True

    for childNode in node.children:

        if childNode is not visited:
            generate(childNode)

   arr.Add(node.id)

The function is called on the root node of the graph, which is 1.

For example:

You are given N = 4, ‘arr’ = [1, 2, 2, 1].The resulting graph will be

altTExt

For node 1, the subtree count is 2, while for node 2 the subtree count is 1. Hence the answer is [2, 1].

Input Format:

The first line of input contains an integer ‘T’ representing the number of test cases.

The first line of each test case contains a single integer ‘N’ representing the size of the array.

The second line of input contains ‘N’ space-separated integers representing the given array.

Output Format:

For each test case, print space-separated integers, representing the number of nodes in each subtree. 

Print a separate line for each test case.

Constraints:

1 <= T <= 10
2 <= N <= 10^6
0 <= arr[i] <= 10^9

Time Limit: 1 sec

Note:

You do not need to print anything. It has already been taken care of. Just implement the function.

Solve now

Approaches

01 Approach

In this approach, we make a few observations,

The array will always be of the size of 2*number of nodes
Whenever a node id reoccurs in the array then there will not be any more children of that node in the array

We can build the tree by pushing the node id in a stack and adding the next node as the children of the top of the stack. Whenever the top of the stack is equal to the current node id in the array, then we will pop the node from the stack

After building the tree we can use DFS to find the nodes in each subtree of the tree

We will create a class TreeNode(id, children), where id is the id of the node and children is the array of child nodes of the current node.

We will create a function getSubTree(root, count), here root is the current node and count the array of subtree counts of the tree.

Algorithm:

In the function getSubTree(root, count)
- Set count[root.val - 1] as 1
- Iterate child through root.children
  - Call getSubTree(child, count)
  - Add count[child.id - 1] to count[root.id - 1]
Initialise an empty stack st
Initialise TreeNode(0) as root
Insert root into st
Set currIndex as 1
While st is not empty
- If arr[currIndex] is equal to top of st
  - Pop the stack
- Otherwise,
  - Initialise newNode as TreeNode(arr[currIndex])
  - Insert newNode into st.top.children
  - Insert newNode in st
- Increase currIndex by 1
Initialise an empty array count of half the length of arr
Call getSubTree(root, count)
Return count

Solve now

02 Approach

In this approach, it can be seen that all the node IDs, between two occurrences of the same node ID in the array, are the subtree nodes of the array, and all nodes are also occurring twice. Hence the number of nodes in subtree of a node ID is half of the distance between the two occurrences of the Node ID

Algorithm:

Initialise an empty map lastFound
Initialise an empty array count
Iterate index and num of the arr
- If num is in lastFound
  - Set count[num - 1] as (index - lastFound[num] + 1) / 2
- Otherwise
  - Set lastFound[num] as index

Return count array

Solve now