AtCoder が提供しているABC（AtCoder Beginner Contest）306 のE問題をC++で解いてみました。ABC306は、2023年6月17日21:00に実施されました。
この回は、コンテスト中にジャッジ遅れが大きく発生したため、unrated となりました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Best Performances（Difficulty : 1268）
問題はリンク先をご覧ください。
multisetを使って解くことができました。AtCoder Problems による Difficulty は 1268 でした。
解答案
数列 A から、降順に K 個の要素を格納するコンテナ s と、それ以外の要素を格納するコンテナ t を分けて管理します。コンテナに求められる要件は以下です。
- （順序を保ったまま）要素の挿入ができる（s, t 共に）。
- （順序を保ったまま）要素の削除ができる（s, t 共に）。
- 最小の値が取り出せる（s に対する要件）。
- 最大の値が取り出せる（t に対する要件）。
少ない計算量で上記の条件を満たすコンテナとして二分探索木を使う set があります。この問題は、重複する要素があるため、mulitset を使います。
multiset については、AOJ ITP2 の 7-D で紹介しています。
C++ プログラム例（ABC306E）
基本的に xi と yi と数列 ai の値に従い、コンテナ s と t を更新していきます。出力する解答 result は、s と t の更新に合わせて、更新します。
プログラムとしての工夫は以下です。
- 配列の添え字は0から始まるため、xi も同じように扱います（13行目）。
- 更新する後の値を変数 next、現在の値を変数 now としました。最後に ai を更新します（29、30、47行目）。
- next の処理をするときは、先に挿入しています。これはコンテナの要素がないときに削除することを防ぐためです。s に next を挿入して、この時点で s の一番小さい値を s から t に移します（32ー35行目）。
- コンテナ s に now が存在する場合は、now を削除して t の一番大きな値を t から s に移します（39-42行目）。s に now が存在しない場合は、t から now を削除します（45行目）。
- 更新毎に出力する変数 result は、32ビット整数に収まらない可能性があるため、64ビット整数にします（27行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n, k, q;
cin >> n >> k >> q;
vector<int> x(q), y(q);
for (int i = 0; i < q; ++i) {
cin >> x[i] >> y[i];
--x[i];
}
vector<int> a(n, 0);
multiset<int> s;
for (int i = 0; i < k; ++i) {
s.insert(0);
}
multiset<int> t;
for (int i = 0; i < n - k; ++i) {
t.insert(0);
}
ll result = 0;
for (int i = 0; i < q; ++i) {
int now = a[x[i]];
int next = y[i];
s.insert(next);
int smallest = *s.begin();
s.erase(s.find(smallest));
t.insert(smallest);
result = result + next - smallest;
if (s.find(now) != s.end()) {
s.erase(s.find(now));
int biggest = *t.rbegin();
s.insert(biggest);
t.erase(t.find(biggest));
result = result + biggest - now;
} else {
t.erase(t.find(now));
}
a[x[i]] = y[i];
cout << result << endl;
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python の ordered set について
Python には、標準で順序付きの set がありません。ネットで探すといくつかの実装を見つけることができますが、記事として取り上げることを見送ります。
最後に
ABC306は、30分以内にD問題を解くことができました。一方、E問題は1時間格闘しましたが、提出に至りませんでした。
解けなかった要因は以下でした。
- multiset を使うことに気づいていた。しかし erase をイテレータではなく値で行っていた。このため、同じ値を持つ要素がすべて削除されていた。
- この問題に気づいてmapに切り替えて実装したところで時間切れ。
- next と now の大きさにより条件文を分けて、ごちゃごちゃしてしまった。
公式解説で、自分のプログラムの問題点を理解することができて学べました。
引き続き ABC の問題を紹介していきます。
