(ABC371)D - 1D Country

はじめに

9/14開催のABC371のD問題です。

問題文

解釈

N個の村があって、$i$番目の村は座標$X_i$で$P_i$人の村人がいる。
以下の計算をして、返却する。

• 整数$L_i,R_i$が与えられるので$L_i \leq X_i \leq R_i$に含まれる村すべての村人を合計

制約

• $1 \leq N,Q \leq 2 \times 10^5$
• $-10^9 \le X_1 < X_2 < \cdots < X_N \le 10^9$
• $1 \le P_i \le 10^9$
• $-10^9 \le L_i \le R_i \le 10^9$

考え方

制約の数値が非常に大きいので、$X_i$を１つずつ見つけたり、加えたりすると恐怖のTLEになるのは目に見えてます。

村をさっと見つけて、合計を少ない計算で出す方法を検討する。

村をさっと見つける => 二分探索
合計を少ない計算　=> 累積和

これは数少ない私の知っているアルゴリズムでなんとかなりそうです。

自分の解答

ABC371_D.py

def binary_search(nums ,x):
  left , right = 0 , len(nums)
  while left < right:
    mid = (left + right) // 2
    if nums[mid][0] < x:
      left = mid + 1
    else:
      right = mid
  return left

N = int(input())
X = list(map(int,input().split()))
P = list(map(int,input().split()))
Q = int(input())

#クエリをまとめて配列にする
queries = [list(map(int, input().split())) for _ in range(Q)] 

#村の座標と人口をまとめて配列にする
villages = list(zip(X, P)) 

#累積和を計算
cumulative_sum = [0] * (N + 1)
for i in range(N):
  cumulative_sum[i + 1] = cumulative_sum[i] + villages[i][1]

#クエリを処理
for L, R in queries:
    left = binary_search(villages, L)
    right = binary_search(villages, R + 1)
    print(cumulative_sum[right] - cumulative_sum[left])

アルゴリズムの勉強をしてきて、初めてちゃんと使えたので、解けた時、脳汁噴出したのがわかりました。

おわりに

勉強したアルゴリズムが身についているのかを確認するため、説明する記事を今後は書いていきたいと思います。