AtCoder が提供しているABC（AtCoder Beginner Contest）321 のD問題をC++とPythonで解いてみました。ABC321は、2023年9月23日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Set Menu（Difficulty : 806）
問題はリンク先をご覧ください。
累積和と二分探索を組み合わせて解くことができました。AtCoder Problems による Difficulty は 806 でした。
解答案
$1 \leqq N,\, M \leqq 2 \times 10^5$ であるため、主菜と副菜のすべての組み合わせの価格を計算することは時間的に間に合いません。
主菜価格（配列 a）をソートして、副菜をひとつ固定します。
- 副菜の価格が決まっているため、ソートした主菜価格に二分探索を用いると、定食価格がセット価格（P）以下になる主菜の種類が分かります。
- 主菜価格の累積和を求めておけば、定食価格の総和の計算が簡単になります。
この工夫により、ソートと合わせて定食価格の総和を求める計算量は、$N \log N + M \log N$ となり、時間的に間に合うようになりました。
C++ プログラム例（ABC321D）
以下は、プログラムの補足です。
- 主菜価格を配列 a に格納してソートします。ソート済の配列 a の累積和 a_sum も求めておきます（20ー25行目）。
- p – b[i] で配列 a を二分探索して、セット価格以下となるメニュー数 c を求めます。n – c 種類の定食がセット価格になります（29行目）。
- 主菜と副菜を組み合わせる価格（30行目）とセット価格（31行目）を解答 result に加えます。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n, m;
ll p;
cin >> n >> m >> p;
vector<ll> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<ll> b(m);
for (int i = 0; i < m; ++i) {
cin >> b[i];
}
sort(a.begin(), a.end());
vector<ll> a_sum(n + 1);
a_sum[0] = 0;
for (int i = 0; i < n; ++i) {
a_sum[i + 1] = a_sum[i] + a[i];
}
ll result = 0;
for (int i = 0; i < m; ++i) {
int c = upper_bound(a.begin(), a.end(), p - b[i]) - a.begin();
result += a_sum[c] + b[i] * c;
result += (n - c) * p;
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC321D）
Python 版も基本的な考え方は同じです。C++ の upper_bound は、bisect.bisect_right に置き換えられます（2、15行目）。以下となります。
"""AtCoder Beginner Contest 321 D"""
import bisect
n, m, p = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
a.sort()
a_sum = [0] * (n + 1)
for i in range(n):
a_sum[i + 1] = a_sum[i] + a[i]
result = 0
for i in range(m):
c = bisect.bisect_right(a, p - b[i])
result += a_sum[c] + b[i] * c
result += (n - c) * p
print(result)こちらも「AC」と判定されました。
最後に
累積和と二分探索が必要な面白い問題でした。
コンテストでは以下の実装ミスがありました。実装中に気付いたのでペナルティを避けることができましたが、時間を失いました。
- 累積和の計算が良くなくて、添え字1個違いミスをした。この記事のプログラムでは直しています。
- 計算式を何度も間違えた。紙に書いて頭を整理してプログラムするべきだった。
- 型を unsigned long long にしていて、p-b[i] が負になる場合に計算が合わない。
このような実装ミスを減らして、他の問題に使える時間を増やしたいです。
引き続き ABC の問題を紹介していきます。
