AtCoder が提供しているABC（AtCoder Beginner Contest）371 D問題をC++とPythonで解いてみました。ABC371は、2024年9月14日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 1D Country（Difficulty : 408）
問題の詳細は、リンク先をご覧ください。
累積和と二分探索を組み合わせます。AtCoder Problems による Difficulty は 408 でした。
解答案
C++ プログラム例（ABC371D）
村人の人口の累積和 $s$ を求めておきます。これにより $[L_i R_i]$ に含まれる人口は、$L_i$ 以上の $X$ のインデックス $iL$ と、$R_i$ より大きな $X$ のインデックス $iR$ を求めれば、$s_{iR} – s_{iL}$ で求めることができます。
プログラムを補足します。
- 村の人口の累積を s を計算しておきます（19ー22行目）。
- クエリに対して
- $L_i$ 以上のインデックス $iL$ を
lower_boundで求めます（29行目）。 - $R_i$ より大きいインデックス $iR$ を
upper_boundで求めます（30行目）。 - 累積和から指定された範囲の村人の人口を求めます（31行目）。
- $L_i$ 以上のインデックス $iL$ を
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n;
cin >> n;
vector<ll> x(n);
for (int i = 0; i < n; ++i) {
cin >> x[i];
}
vector<ll> p(n);
for (int i = 0; i < n; ++i) {
cin >> p[i];
}
vector<ll> s(n + 1);
for (int i = 0; i < n; ++i) {
s[i + 1] = s[i] + p[i];
}
int q;
cin >> q;
for (int i = 0; i < q; ++i) {
ll L, R;
cin >> L >> R;
auto iL = lower_bound(x.begin(), x.end(), L) - x.begin();
auto iR = upper_bound(x.begin(), x.end(), R) - x.begin();
cout << s[iR] - s[iL] << endl;
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC371D）
Python版も基本的な考え方は同じです。C++ の lower_bound は bisect_left に upper_bound は bisect_right に相当します。以下がプログラムです。全体がすっきりと書けています。
"""AtCoder Beginner Contest 371 D"""
import bisect
n = int(input())
x = list(map(int, input().split()))
p = list(map(int, input().split()))
s = [0] * (n + 1)
for i in range(n):
s[i + 1] = s[i] + p[i]
q = int(input())
for i in range(q):
L, R = map(int, input().split())
iL = bisect.bisect_left(x, L)
iR = bisect.bisect_right(x, R)
print(s[iR] - s[iL])こちらも「AC」と判定されました。
最後に
累積和と二分探索を使う良い問題だと思います。
引き続き ABC の問題を紹介していきます。
