coding-problems/problems/codechef/SHROUTE/sol.py at master · NicoKNL/coding-problems · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
marker = 1000000000


def solve():
    N, M = [int(x) for x in str(input()).split(" ")]
    A = [0]
    A.extend([int(x) for x in str(input()).split(" ")])
    B = [int(x) for x in str(input()).split(" ")]
    C = [marker] * (N + 1)

    # Dist from left
    has_train = False
    train_pos = -1
    for i in range(1, len(A)):
        if A[i] == 1:
            has_train = True
            train_pos = i
            C[i] = 0
        elif has_train:
            C[i] = i - train_pos
            i += 1

    # Dist from right
    has_train = False
    train_pos = -1
    for i in range(len(A) - 1, 0, -1):
        if A[i] == 2:
            has_train = True
            train_pos = i
            C[i] = 0
        elif has_train:
            C[i] = min(C[i], train_pos - i)
            i += 1

    # Set magic station
    C[1] = 0

    # Build result
    result = ""
    for i in range(len(B)):
        if i > 0:
            result += " "

        destination = B[i]
        if C[destination] == marker:
            result += "-1"
        else:
            result += str(C[destination])

    print(result)


if __name__ == "__main__":
    T = int(input())

    while T:
        solve()
        T -= 1