AtCoder が提供しているABC（AtCoder Beginner Contest）281 のE問題をC++で解いてみました。ABC281は、2022年12月10日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
E問題 Least Elements（Difficulty : 1334）
問題はリンク先をご覧ください。
multisetを使って解くことができました。AtCoder Problems による Difficulty は 1334 でした。
解答案
数列 A から連続する M 個の要素を取り出し、昇順に K 個の要素を格納するコンテナ s と、それ以外の要素を格納するコンテナ t を分けて管理します。コンテナに求められる要件は以下です。
- （順序を保ったまま）要素の挿入ができる（s, t 共に）。
- （順序を保ったまま）要素の削除ができる（s, t 共に）。
- 最大の値が取り出せる（s に対する要件）。
- 最小の値が取り出せる（t に対する要件）。
少ない計算量で上記の条件を満たすコンテナとして二分探索木を使う set があります。この問題は、重複する要素があるため、mulitset を使います。
multiset については、AOJ ITP2 の 7-D で紹介しています。
C++ プログラム例（ABC281E）
数列 ai の値に従い、コンテナ s と t を更新していきます。出力する解答 result は、s と t の更新に合わせて、更新します。
プログラムとしての処理は以下です。
- 最初の回答は、ループ前に計算して出力しておきます（24、31行目）。
- 更新する後の値を変数 next、現在の値を変数 now としました（33、34行目）。
- next の処理をするときは、先に挿入しています。これはコンテナの要素がないときに削除することを防ぐためです。s に next を挿入して、この時点で s の一番大きい値を s から t に移します（36ー39行目）。
- コンテナ s に now が存在する場合は、now を削除して t の一番小さな値を t から s に移します（43ー46行目）。s に now が存在しない場合は、t から now を削除します（49行目）。
- 更新毎に出力する変数 result は、32ビット整数に収まらない可能性があるため、64ビット整数にします（20行目）。
以下が、C＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
int main()
{
int n, m, k;
cin >> n >> m >> k;
vector<ll> a(n);
vector<ll> b(m);
for (int i = 0; i < n; ++i) {
cin >> a[i];
if (i < m) {
b[i] = a[i];
}
}
sort(b.begin(), b.end());
ll result = 0;
multiset<int> s;
for (int i = 0; i < k; ++i) {
s.insert(b[i]);
result += b[i];
}
multiset<int> t;
for (int i = k; i < m; ++i) {
t.insert(b[i]);
}
cout << result << " ";
for (int i = m; i < n; ++i) {
int now = a[i - m];
int next = a[i];
s.insert(next);
int biggest = *s.rbegin();
s.erase(s.find(biggest));
t.insert(biggest);
result = result + next - biggest;
if (s.find(now) != s.end()) {
s.erase(s.find(now));
int smallest = *t.begin();
s.insert(smallest);
t.erase(t.find(smallest));
result = result + smallest - now;
} else {
t.erase(t.find(now));
}
cout << result << " ";
}
cout << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python の ordered set について
Python には、標準で順序付きの set がありません。ネットで探すといくつかの実装を見つけることができますが、記事として取り上げることを見送ります。
この問題と似た ABC306E は、Difficulty 1268 でした。素の Python で一工夫必要な問題の Difficulty は高くなる傾向があるのかもしれません。
最後に
ABC306E を解くことで、この問題も同じ構造を持っていることに気づきました。記事にする時間は前後しますが、紹介することにしました（記事とプログラムの構成が似ています）。
引き続き ABC の問題を紹介していきます。
