0 はじめに

0-1 こんにちは

こんにちは。
結構前にABCの解説を投げてた人です。
久し振りに、ABC解説以外に投稿するネタが降って来たので書きます。

0-2 この記事について

先日のABCのD問題「Bank」―C++だとsetで殴れる問題―をPythonでデータ構造を使って解いてみます。
確認していないので分かりませんが、既にこのような記事を書いている方がいたらごめんなさい。

1 本題

1-1 問題の概要

N人の人と受付があります。
次のようなクエリがたくさん降ってくるので随時処理して答えて下さい。

1 : 受付に呼ばれていない人の中で最小番号の人が呼ばれる
2 x : 人x(既に呼ばれている)が初めて受付へ行く
3 : 既に受付に呼ばれているが行っていない人のうち最小番号の人が呼ばれる
　　その番号を出力せよ

要は、最小値の取得と任意の要素の削除が高速に出来ればよいです。

1-2 C++だとどのように殴れるか

C++だとstd::setを用いて殴るだけです。
std::setはつよくなった「集合を扱う型」みたいな感じで、こんなことが出来ます :

(Sはset型の変数、nは引数とします)

できること	記述	計算量
要素の挿入	`S.insert(n)`	$O(\log N)$
要素の削除	`S.erase(n)`	$O(\log N)$
要素の存在判定	`S.count(n)`	$O(\log N)$
最小値の取得	`*begin(S)`	$O(1)$
最大値の取得	`*rbegin(S)`	$O(1)$

つよいですね。

要素の挿入/削除などに$O(\log N)$掛かるのは少し気になってしまいますが(?)、それでいて最大値/最小値も$O(1)$で取得できるのはつよいです。

実装例を貼っておきます。

C++での実装例

#include <bits/stdc++.h>
using namespace std;
using ll=long long;
#define rep(i,N,M) for(int i=N; i<M; i++)

int main(){
  int N,Q; cin>>N>>Q;
  
  set<int> called;
  int uncalled=0;
  rep(i,0,Q){
    int n; cin>>n;
    if(n==1){
      called.insert(uncalled);
      uncalled++;
    }else if(n==2){
      int x; cin>>x; x--;
      called.erase(x);
    }else{
      cout<<*called.begin()+1<<endl;
    }
  }
}

受付に呼ばれる人の番号は絶対に$1,2,3,...$となっている為、呼ばれた度に$1$増やせばよく、
setで管理する必要があるのは呼ばれた人の番号のみです。
上のように殴れます。
計算量は$O(Q \log N)$で十分高速ですね。(実行時間 : 717ms)

1-3 Pythonでは？

残念ながら、C++のstd::setに相当するものがPythonの標準ライブラリにはありません！！！

え、PythonのsetはC++だと何なんだって？
それがね、unordered_setなんですね....(哀)

最小値の取得と要素の削除が両方高速に行えるデータ構造はPythonにはありません。

1-4 Pythonではどうできるのか

そこで、heapqとsetを併用することを考えます。(heapqベースで書くことを前提とします)

この問題を解く上で足枷となっているのは2のクエリです。
1は、C++のところで述べた通りで、簡単です。
3も、最小値を取得するだけなのでheapqを用いるなら簡単です。

うーん、
あ！天啓が降って来た━━━━(ﾟ∀ﾟ)━━━━!!

こんな工夫をしてみましょう。

heapqで既に呼ばれた人(受付に行ったとは限らない)を、setで既に受付に行った人を管理する。
2のクエリでは受付に行く人をsetに追加する。
3のクエリでは、最小値を取得&heappopを、setに含まれない要素が出て来るまで繰り返し、setに含まれない最小の値が答え。

こうすることで、setとheapqだけで殴ることができます。
heappopの回数も多くても$N$回なので計算量は$O(Q \log N)$で十分高速です。
賢くないですか？

実装してみました。

Pythonでの実装

from heapq import heapify, heappop, heappush

N,Q=map(int,input().split())
uncalled=0
called=[]; heapify(called)
gone=[0]*N;
for i in range(Q):
  query=list(map(int,input().split()))
  c=query[0]
  if c==1:
    heappush(called,uncalled)
    uncalled+=1
  elif c==2:
    n=query[1]; n-=1
    gone[n]=1
  else:
    while True:
      mn=called[0]
      if not gone[mn]:
        print(mn+1)
        break
      heappop(called)

特にクエリ3の芸術点高いと思います。
実行時間は328msでした。はやい！！(入力高速化しましたが)

2 おわりに

もっと速い別解があると知って虚無の顔をしています。
こんな感じのことをしてPythonでも解けるよ～みたいな問題にはABC253-Cがありましたね。

いいね頂けると幸いです。
お読み頂きありがとうございました。