AtCoder が提供しているABC（AtCoder Beginner Contest）331 のE問題をC++とPythonで解いてみました。ABC331は、2023年12月2日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Set Meal（Difficulty : 1018）
問題はリンク先をご覧ください。
食べ合わせが悪い組み合わせをsetコンテナに格納して処理します。AtCoder Problems による Difficulty は 1018 でした。
解答案
主菜の種類 $N$ と副菜の種類 $M$ の制約は、$1 \leqq N, M \leqq 10^5$ であるため、単純にすべての組み合わせを調べることはできません。
この問題は、主菜 a[i] と副菜 b[j] の値段を組み合わせた最大値を求めるため、主菜 a[i] を固定して、副菜 b[j] を大きな値から確認します。ある b[j] が a[i] と食べ合わせが悪くなければ、これより小さい b[j’] について調べる必要はありません。
食べ合わせの悪い組み合わせ $L$ は、最大でも $10^5$ 種類です。set コンテナをうまく使えば、間に合います。
C++ プログラム例（ABC331E）
プログラムの補足です。
- 値とインデックスを合わせた配列 bs を降順にソートする（19、21行目）。
- 主菜毎に食べ合わせが悪い副菜を set コンテナに登録する（29行目）。
- 主菜 a[i] に対して、値段が大きい方から副菜 bs[j] を取り出して、食べ合わせが悪くなければ、ループを抜けて、最大値 result を更新します。
配列 bs[j] を調べるループ（35ー42行目）が、$N \times M$ 回のループになっているように見えます。この場合は、食べ合わせが悪い組み合わせの回数 $L$ に対して、最大 $L + N$ 回しかループは実行されないため間に合います。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
#define INF (1LL << 60)
int main()
{
int n, m, L;
cin >> n >> m >> L;
vector<ll> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
vector<ll> b(m);
vector<pair<ll, int>> bs(m);
for (int i = 0; i < m; ++i) {
cin >> b[i];
bs[i] = make_pair(b[i], i);
}
sort(bs.begin(), bs.end(), greater<pair<ll, int>>());
vector<set<int>> p(n);
for (int i = 0; i < L; ++i) {
int c, d;
cin >> c >> d;
--c;
--d;
p[c].insert(d);
}
ll result = 0;
for (int i = 0; i < n; ++i) {
ll b_max = -INF;
for (int j = 0; j < m; ++j) {
ll price = bs[j].first;
int index = bs[j].second;
if (p[i].find(index) == p[i].end()) {
b_max = price;
break;
}
}
result = max(result, a[i] + b_max);
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC331E）
Python 版も基本的な考え方は同じです。以下となります。
"""AtCoder Beginner Contest 331 E"""
INF = 1 << 60
n, m, L = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
bs = []
for i in range(m):
bs.append((b[i], i))
bs.sort(reverse=True)
p = [set() for _ in range(n)]
for i in range(L):
c, d = map(int, input().split())
c -= 1
d -= 1
p[c].add(d)
result = 0
for i in range(n):
b_max = -INF
for j in range(m):
price = bs[j][0]
index = bs[j][1]
if index not in p[i]:
b_max = price
break
result = max(result, a[i] + b_max)
print(result)こちらも「AC」と判定されました。
最後に
この問題の Difficulty は、1018でした。自分のレートから考えるとやや難しい問題ですが、setコンテナに慣れてきたのか、手際よく解くことができました。このような問題を増やしていきたいです。
引き続き ABC の問題を紹介していきます。
