AtCoder が提供しているABC（AtCoder Beginner Contest）334 のD問題をC++とPythonで解いてみました。ABC334は、2023年12月23日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Reindeer and Sleigh（Difficulty : 602）
問題はリンク先をご覧ください。
ABC334 D問題 Reindeer and Sleigh
ソートして累積和を求めて二分探索です。AtCoder Problems による Difficulty は 602 でした。
解答案
ソリを引くためのトナカイの数をソートしておきます。これで、小さい方から加えていくことにより、ひとつのクエリに答えることができます。ただし、計算量が $O(NQ) = 4 \times 10^{10}$ となり間に合いません。
前準備として、トナカイの数の累積和を求めておきます。（トナカイの数）X 以下となる累積和の要素のインデックスは、二分探索で求めることができます。この場合、ひとつのクエリを $O(\log N)$ で処理することができます。結果的に間に合います。
C++ プログラム例（ABC334D）
上で述べたように、ソートして（14行目）、累積和を求めて（16ー20行目）、二分探索で解を求めます（25行目）。注意としては、累積和の配列 s は、s[0] ＝ 0 を持つ必要があるため（※）、求めた要素数から1を引く必要があります。
※ [left , fight) の累積和を s[right] – s[left] で求めることができます。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n, q;
cin >> n >> q;
vector<ll> r(n);
for (int i = 0; i < n; ++i) {
cin >> r[i];
}
sort(r.begin(), r.end());
vector<ll> s(n + 1);
s[0] = 0;
for (int i = 0; i < n; ++i) {
s[i + 1] = s[i] + r[i];
}
for (int i = 0; i < q; ++i) {
ll x;
cin >> x;
int result = upper_bound(s.begin(), s.end(), x) - s.begin() - 1;
cout << result << endl;
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC334D）
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 334 D"""
import bisect
n, q = map(int, input().split())
r = list(map(int, input().split()))
r.sort()
s = [0] * (n + 1)
s[0] = 0
for i in range(n):
s[i + 1] = s[i] + r[i]
for i in range(q):
x = int(input())
result = bisect.bisect_right(s, x) - 1
print(result)こちらも「AC」と判定されました。
最後に
この問題は、累積和と二分探索という2つの手法を組み合わせている教育効果が高い問題だと思います。
引き続き ABC の問題を紹介していきます。
