Problem p02599 (python)

Fixed submission:
from sys import stdin
#import numpy as np

readline = stdin.readline

N, Q = map(int, readline().split()) 
C = list(map(int, readline().split())) 
S = [list(map(int, readline().split())) for _ in range(Q)] 
#print(S)
def main(N, Q, C, S):    
  # Binary Indexed Tree 
  # https://juppy.hatenablog.com/entry/2018/11/17/%E8%9F%BB%E6%9C%AC_python_Binary_Indexed_Tree_%E7%AB%B6%E6%8A%80%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3%E3%82%B0
  b = [0 for _ in range(N+1)]
  def BIT_query(idx): #idxまでの和(1からスタート)
    res_sum = 0
    while idx > 0:
      res_sum += b[idx]
      idx -= idx&(-idx) #一番下の位を減らす（例：13=1011 → 1010 → 1000） 
    return res_sum
  def BIT_update(idx,x): #idxをx増やした時
    while idx <= N:
      b[idx] += x
      idx += idx&(-idx) #一番下の位を繰り上げる（例：13=1011 → 1100 → 10000 ※N=16のとき） 
    return

  loc = [0 for _ in range(N)] #各種類の玉のうちの最右端の所在地
  S = [(S[i], i) for i in range(Q)] #クエリ
  S.sort(key = lambda s: s[0][1]) #区間の右端でソート
  #S = np.array(S, dtype = np.int64)
  #print(S)
  ans = [0 for _ in range(Q)]
  k = 0
  for j in range(N):
    if loc[C[j]-1] > 0: #玉の更新
      BIT_update(loc[C[j]-1], -1)
    loc[C[j]-1] = j+1
    BIT_update(j+1, 1)  
    
    if S[k][0][1] == j+1:
      t = BIT_query(j+1)
    while k < len(S) and S[k][0][1] == j+1: #クエリへの返答
      ans[S[k][1]] = t - BIT_query(S[k][0][0]-1)
      k += 1
    if k == len(S):
      break
  return ans

ans = main(N, Q, C, S)
print(*ans, sep='\n')

Buggy submission (verdict: Runtime Error):
from heapq import heapify,heappush,heappop
from sys import stdin
input = stdin.readline

def arr_4(x):
  return [x,x,x,x]

H, W, K = map(int, input().split())
INF = 10 * H * W

x1, y1, x2, y2 = map(int, input().split())
start = (x1-1 ,y1-1, 0)
goal = (x2-1, y2-1, 0)

G = [list(map(arr_4,input().rstrip())) for _ in range(H)]
#G[i][j][k]: (i,j)座標、k方向(0-3)

def dijkstra(G, start, goal):# G:格子、G[i][j][l] = 0 なら通過可、1ならコスト+1 
  H = len(G)
  W = len(G[0])
  #Dir = 4 方向
  task = [] 
  min_cost = [[[(INF,0) for _ in range(4)] for _ in range(W)] for _ in range(H)]
  cost = [[[(INF,0) for _ in range(4)] for _ in range(W)] for _ in range(H)]
  #prev_points = [[[(-2,-2) for _ in range(4)] for _ in range(W)] for _ in range(H)]
  for i in range(4):
    heappush(task, ((1, 0), (start[0], start[1], i), start)) 
    
  #（開始）頂点、コスト、前の点
  # !!コストは（加速回数、慣性で進んだ距離）で規定!! #
  # !!最初の点番号が0でないときは、真ん中の0を変更する!! #
  
  flag = 0
  while task:
    while task:
      c, p, prev = heappop(task)    
      if not G[p[0]][p[1]][p[2]] == "@":
        break  
      if len(task) == 0:
        flag = 1
    if flag == 1:
      break
    #print(task)
    G[p[0]][p[1]][p[2]] = "@"
    min_cost[p[0]][p[1]][p[2]] = c
    #prev_points[p[0]][p[1]] = prev
    
    if p[0] == goal[0] and p[1] == goal[1]:
      return min_cost#, prev_points)
  
    #次の頂点を列挙
    q = []
    if p[0] > 0 and p[2] == 0:#南
      q.append((p[0]-1,p[1],p[2])) 
    if p[0] < H-1 and p[2] == 1:#北
      q.append((p[0]+1,p[1],p[2])) 
    if p[1] > 0 and p[2] == 2:#西
      q.append((p[0],p[1]-1,p[2])) 
    if p[1] < W-1 and p[2] == 3:#東
      q.append((p[0],p[1]+1,p[2]))   
    for i in range(4):
      if not i == p[2]:
        q.append((p[0],p[1],i)) 
    
    for j in q: #制約が入るときはここに追加
      if G[j[0]][j[1]][j[2]] == "@": continue
      
      if p[2] == j[2]:#向き変更なし
        if c[1] == K:
          cost_next = (c[0]+1,1)
        else:
          cost_next = (c[0],c[1]+1)
      else:#向き変更
        cost_next = (c[0]+1,0)
      
      if cost[j[0]][j[1]][j[2]] > cost_next:
        cost[j[0]][j[1]][j[2]] = cost_next
        heappush(task, (cost_next, j, p))
        #print(cost_next, j, p)  
      #print(task)
      
    del q
    
  return min_cost#, prev_points)

min_cost = dijkstra(G, start, goal)

#print(min_cost)

c = min(min_cost[x2-1][y2-1])[0]
if x1 == x2 and y1 == y2:
  c = 0
print(c if c < INF else -1)

Buggy submission (verdict: Runtime Error):
from sys import stdin
#import numpy as np

readline = stdin.readline

N, Q = map(int, readline().split())
C = list(map(int, readline().split()))
S = [list(map(int, readline().split())) for _ in range(Q)]

def main(N, Q, C, S):    
  # Binary Indexed Tree 
  # https://juppy.hatenablog.com/entry/2018/11/17/%E8%9F%BB%E6%9C%AC_python_Binary_Indexed_Tree_%E7%AB%B6%E6%8A%80%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3%E3%82%B0
  b = [0 for _ in range(N+1)]
  def BIT_query(idx): #idxまでの和(1からスタート)
    res_sum = 0
    while idx > 0:
      res_sum += b[idx]
      idx -= idx&(-idx) #一番下の位を減らす（例：13=1011 → 1010 → 1000） 
    return res_sum
  def BIT_update(idx,x): #idxをx増やした時
    while idx <= N:
      b[idx] += x
      idx += idx&(-idx) #一番下の位を繰り上げる（例：13=1011 → 1100 → 10000 ※N=16のとき） 
    return

  loc = [0 for _ in range(N)] #各種類の玉のうちの最右端の所在地
  S = [[S[i][1], S[i][0], i] for i in range(Q)] #クエリ
  S.sort() #区間の右端でソート
  #S = np.array(S, dtype = np.int64)
  #print(S)
  ans = [0 for _ in range(N)]
  k = 0
  for j in range(N):
    if loc[C[j]-1] > 0: #玉の更新
      BIT_update(loc[C[j]-1], -1)
    loc[C[j]-1] = j+1
    BIT_update(j+1, 1)  
    
    if S[k][0] == j+1:
      t = BIT_query(j+1)
    while k < len(S) and S[k][0] == j+1: #クエリへの返答
      ans[S[k][2]] = t - BIT_query(S[k][1]-1)
      k += 1
    if k == len(S):
      break
  return ans

ans = main(N, Q, C, S)
for a in ans:
  print(int(a))

Buggy submission (verdict: Runtime Error):
from sys import stdin
import numpy as np

readline = stdin.readline

N, Q = map(int, readline().split())
C = np.array(list(map(int, readline().split())), dtype=np.int64)
S = np.array([list(map(int, readline().split())) for _ in range(Q)], dtype=np.int64).reshape(Q,2)

def main(N, Q, C, S):    
  # Binary Indexed Tree 
  # https://juppy.hatenablog.com/entry/2018/11/17/%E8%9F%BB%E6%9C%AC_python_Binary_Indexed_Tree_%E7%AB%B6%E6%8A%80%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3%E3%82%B0
  b = np.zeros(N+1, dtype = np.int64)
  def BIT_query(idx): #idxまでの和(1からスタート)
    res_sum = 0
    while idx > 0:
      res_sum += b[idx]
      idx -= idx&(-idx) #一番下の位を減らす（例：13=1011 → 1010 → 1000） 
    return res_sum
  def BIT_update(idx,x): #idxをx増やした時
    while idx <= N:
      b[idx] += x
      idx += idx&(-idx) #一番下の位を繰り上げる（例：13=1011 → 1100 → 10000 ※N=16のとき） 
    return

  loc = np.zeros(N, dtype = np.int64) #各種類の玉のうちの最右端の所在地
  S = [[S[i,1], S[i,0], i] for i in range(Q)] #クエリ
  S.sort() #区間の右端でソート
  S = np.array(S, dtype = np.int64)
  #print(S)
  ans = np.zeros(Q, dtype = np.int64)
  k = 0
  for j in range(N):
    if loc[C[j]-1] > 0: #玉の更新
      BIT_update(loc[C[j]-1], -1)
    loc[C[j]-1] = j+1
    BIT_update(j+1, 1)  
    
    if S[k,0] == j+1:
      t = BIT_query(j+1)
    while k < len(S) and S[k][0] == j+1: #クエリへの返答
      ans[S[k,2]] = t - BIT_query(S[k,1]-1)
      k += 1
    if k == len(S):
      break
  return ans



ans = main(N, Q, C, S)
for a in ans:
  print(int(a))