PythonでJOI難易度6を埋める #20

D - JOI国のお散歩事情 (Walking in JOI Kingdom)

それぞれの国民が他の国民と会い，おしゃべりを始める点を調べる．番目の国民をとする．が東向きに進むとき，よりも東にいて，に一番近い西に進む国民をとすると，十分な時間があれば，二人はとなる．ただし，東に進む国民は，西側から調べていくと更新がうまくできないので，東側からもう一度調べる．もちろん，自分の進行方向に自分と逆向きに進む国民がいなければ，その国民は最初の座標から時間だけ動く．
全ての国民がおしゃべりを始める点が分かれば，与えられた時間で到達可能ならその座標を，到達不可能なら，最初の座標+を出力すればよい．

from bisect import bisect_left
from sys import stdin
input = stdin.readline

n, t, q = map(int, input().strip().split())
nation = [list(map(int, input().strip().split())) for _ in range(n)]
query = [int(input().strip()) for _ in range(q)]

index_left = []
index_right = []

for i in range(n):
    if nation[i][1] == 1:
        index_right.append(i)
    else:
        index_left.append(i)

state = []
length_l = len(index_left)
length_r = len(index_right)
inf = float('inf')

for i in range(n):
    if i == 0:
        if nation[i][1] == 1:
            index = bisect_left(index_left, i)
            if index == length_l:
                state.append(inf)
            else:
                ave = (nation[index_left[index]][0] + nation[i][0]) // 2
                state.append(ave)
        if nation[i][1] == 2:
            index = bisect_left(index_right, i)
            if index == 0:
                state.append(-inf)
            else:
                ave = (nation[index_right[index - 1]][0] + nation[i][0]) // 2
                state.append(ave)
    if i != 0:
        if nation[i][1] == 1:
            index = bisect_left(index_left, i)
            if index == length_l:
                state.append(inf)
            else:
                ave = (nation[index_left[index]][0] + nation[i][0]) // 2
                state.append(ave)
        if nation[i][1] == 2:
            if nation[i][1] == nation[i - 1][1]:
                state.append(state[-1])
                continue
            index = bisect_left(index_right, i)
            if index == 0:
                state.append(-inf)
            else:
                ave = (nation[index_right[index - 1]][0] + nation[i][0]) // 2
                state.append(ave)

for i in reversed(range(n - 1)):
    if nation[i][1] == nation[i + 1][1] == 1:
        state[i] = state[i + 1]

for i in range(q):
    if state[query[i] - 1] == inf:
        print(nation[query[i] - 1][0] + t)
    elif state[query[i] - 1] == inf:
        print(nation[query[i] - 1][0] - t)
    elif abs(state[query[i] - 1] - nation[query[i] - 1][0]) <= t:
        print(state[query[i] - 1])
    else:
        if nation[query[i] - 1][1] == 1:
            print(nation[query[i] - 1][0] + t)
        if nation[query[i]  -1][1] == 2:
            print(nation[query[i] - 1][0] - t)