Next Element
You are given a sequence of N elements, indexed from 0 to N-1. The 'next' element for each element i is specified by an array next_idx, where next_idx[i] is the index of the next element. This forms a structure similar to a linked list, possibly with cycles. You are also given a starting index start_node. Your task is to find the index of the element that directly follows start_node in this sequence.
If start_node is the last element in its sequence (meaning next_idx[start_node] is -1), output -1.
Input
The first line contains an integer N, the number of elements. The second line contains N integers, the next_idx array, where next_idx[i] is the index of the next element after element i. A value of -1 indicates no next element.
The third line contains the integer start_node, the index of the starting element.
Output
Print a single integer: the index of the element following start_node, or -1 if start_node has no next element.
Constraints
1 <= N <= 1000 -1 <= next_idx[i] < N for all i 0 <= start_node < N
Example 1
5 1 2 3 4 -1 0
1
Starting at index 0, the next index is next_idx[0] which is 1.
Example 2
3 -1 0 1 1
0
Starting at index 1, the next index is next_idx[1] which is 0.
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 2 3 4 -1 0
3 -1 0 1 1