Find Node at Fixed Depth
You are given a structure represented by arrays. There are N nodes, indexed from 0 to N-1. An array next_nodes of size N is given, where next_nodes[i] is the index of the node that follows node i. Node with index -1 signifies the end of a path.
You are also given a starting node start_node and a depth d. Your task is to find the index of the node that is exactly d steps away from start_node. If the path from start_node ends before reaching depth d, or if d is negative, output -1. Assume the structure does not contain cycles.
Input
The first line contains two integers N and d (1 <= N <= 1000, -1000 <= d <= 1000).
The second line contains N integers representing the next_nodes array. next_nodes[i] is between -1 and N-1.
The third line contains the start_node index (0 <= start_node < N).
Output
Print the index of the node at depth d from start_node. If the path is too short or d is negative, print -1.
Constraints
1 <= N <= 1000 -1000 <= d <= 1000 -1 <= next_nodes[i] < N for all i 0 <= start_node < N
Example 1
5 2 1 2 3 4 -1 0
2
From node 0, take 2 steps: 0 -> 1 -> 2. The node at depth 2 is 2.
Example 2
5 3 1 2 -1 4 0 0
-1
From node 0, take 3 steps: 0 -> 1 -> 2 -> -1. The path ends before depth 3, so output -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 2 1 2 3 4 -1 0
5 3 1 2 -1 4 0 0