AtCoder が提供しているABC（AtCoder Beginner Contest）376 E問題をC++とPythonで解いてみました。ABC376は、2024年10月19日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Max × Sum（Difficulty : 1063）
問題の詳細は、リンク先をご覧ください。
優先度付きキューを使用して、大きい値を取り出し、それを合計から減算します。AtCoder Problems による Difficulty は 1063 でした。
解答案
$A_i$ と $B_i$ を組にしてソートします。先頭の要素から $K$ 個選ぶと、求める値の初期値は、
$$A_{K-1} \times \sum^{k-1}_{i=0} B_i$$
となります（0から要素をカウントしています）。
次に $K+1$ 番目の要素を考えます。$B_i$ の和から、一番大きな要素を引いて、$B_K$ を加えた結果と $A_K$ との積が求めた値より小さければ求める値をこの値で更新します。この操作を繰り返して、出力する値を求めます。
$B_i$ の和について、以下を繰り返します。
- $K$ 個の $B_i$ から最大の要素を選び、その要素の値を引く。
- 次の $B_{i+1}$ を加える
この用途のために優先度付きキュー（priority_queue）が使えます。
C++ プログラム例（ABC376E）
上で説明した操作を実装します。
- $A_i$ と $B_i$ を組にしてソートします（22ー26行目）。
- 最初の $K$ 個について式の値を求めます（28ー35行目）。
- $B_i$ の和を変数
b_sumに格納します。 - $B_i$ の最大値を priority_queue
b_maxに格納します。C++の priority_queue は、何も指定しないと最大の値から取り出せます。
- $B_i$ の和を変数
- この式の値を更新します（36ー42行目）。
以下が、C++プログラムです。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int t;
cin >> t;
for (int i = 0; i < t; ++i) {
int n;
int k;
cin >> n >> k;
vector<int> a(n);
for (int j = 0; j < n; ++j) {
cin >> a[j];
}
vector<int> b(n);
for (int j = 0; j < n; ++j) {
cin >> b[j];
}
vector<pair<int, int>> p(n);
for (int j = 0; j < n; ++j) {
p[j] = make_pair(a[j], b[j]);
}
sort(p.begin(), p.end());
priority_queue<int> b_max;
ll b_sum = 0;
ll result = 0;
for (int j = 0; j < k; ++j) {
b_sum += p[j].second;
b_max.push(p[j].second);
}
result = p[k - 1].first * b_sum;
for (int j = k; j < n; ++j) {
b_sum -= b_max.top();
b_max.pop();
b_sum += p[j].second;
b_max.push(p[j].second);
result = min(result, p[j].first * b_sum);
}
cout << result << endl;
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC376E）
Python版も基本的な考え方は同じです。優先度付きキューは、heapq モジュールを使いました。C++の場合と異なり、小さい値を優先して返すため、マイナスの値を push しています（20、23、25行目）。以下がプログラムです。
"""AtCoder Beginner Contest 376 E"""
import heapq
t = int(input())
for i in range(t):
n, k = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
p = []
for j in range(n):
p.append((a[j], b[j]))
p.sort()
b_max = []
heapq.heapify(b_max)
b_sum = 0
result = 0
for j in range(k):
b_sum += p[j][1]
heapq.heappush(b_max, -p[j][1])
result = p[k - 1][0] * b_sum
for j in range(k, n):
b_sum -= -heapq.heappop(b_max)
b_sum += p[j][1]
heapq.heappush(b_max, -p[j][1])
result = min(result, p[j][0] * b_sum)
print(result)こちらも「AC」と判定されました。
最後に
自分にとって難易度が高めの問題でしたが、コンテスト中に解くことができました。このような問題が解けるとコンテスト参加を続けることができます。
引き続き ABC の問題を紹介していきます。
