Count Nodes at Depth K
A tree is given as a parent array, where parent[i] is the parent of node i. The root node has parent[i] == -1. You are also given an integer k. Your task is to count the number of nodes that are exactly at depth k. The root node is considered to be at depth 0.
Input
The first line contains an integer n (1 <= n <= 1000), the number of nodes in the tree. The second line contains n integers representing the parent array. The i-th integer is parent[i]. The third line contains an integer k (0 <= k < n).
Output
Print a single integer, the count of nodes at depth k.
Constraints
1 <= n <= 1000
0 <= k < n
parent[i] is either -1 or an integer j such that 0 <= j < n and j != i.
Exactly one node has parent[i] == -1.
Example 1
5 -1 0 0 1 1 2
2
Nodes 3 and 4 are children of node 1, which is a child of node 0 (root). So nodes 3 and 4 are at depth 2.
Example 2
7 -1 0 0 1 1 2 2 1
2
Nodes 1 and 2 are direct children of the root (node 0), so they are at depth 1.
Hints
# Read the input from stdin and print the answer.
import sys
def main():
data = sys.stdin.read().split()
# your code here
main()
5 -1 0 0 1 1 2
7 -1 0 0 1 1 2 2 1