輪郭をなぞるだけのブログ

浅学菲才のためにおそらく嘘も多い

個数が最小のものを高速に取り出す

Codeforces でこの操作が必要な問題があり、それにかなりの時間を取られていたら終わってしまったので、メモしておく。
問題の説明などを書く気力がないので、操作方法だけ記す。
問題のリンク : Problem - D - Codeforces

例えば、1 回の操作で 0 から x-1 のいずれかの個数を 1 つ増やし、それを n 回繰り返すとする。
操作するたびに、個数が最小のものを高速に取り出したい。
操作ごとに個数を更新してソートするのが間に合えば簡単だが、これだと計算量は O(nx \log x) となり、制約によっては間に合わない。

set を使えば、更新が O(\log x) で行える。
ただ、このとき set にどういう情報を格納するかという問題がある。
個数が最小のものを取り出したいので、(個数、数) という情報を入れておきたいが、これだけだと個数の更新がうまくいかない。
これを解決するために、個数更新時の情報を持っておく。
ある数の更新するときは、直前の情報を set から取り出して、個数を更新して set に入れる。
これで計算量が O(n \log x) に改善された。

ソースコードは以下の通り。
コンテスト中に絞り出したものの類型なので、粗削りである。
デバッグもしていない。
また書き直すかもしれない。

#include <iostream>
#include <vector>
#include <set>
using namespace std;

typedef pair<int, int> P;
set<P> s;           // (個数, 数) の集合
vector<int> cnt;    // 個数更新時の (直前の) 情報

void init(int x){
    for(int i = 0; i < x; ++i)  s.insert(P(0, i));
    cnt.assign(x, 0);
}

int get_value(int m){
    // m の個数を更新する
    s.erase(P(cnt[m]++, m));
    s.insert(P(cnt[m], m));
    
    // 個数が最小のものの数を返す
    auto it = s.begin();
    return (*it).second;
}

他の人が書いてるの見たけど、

vector<set<int>> s;

の方がいいかなあ。

問題の内容を理解して、適切なデータ構造をすぐに選べるようになろうね…