はじめに
ABC 194 E 問題 Mex Min を解くために考えたこと、ACできるPython3(PyPy3)コードを紹介します。
E.Mex Min
問題ページ
難易度 ： 緑色 1088
考察
データ構造の選択
$mex$ : 区間に含まれない & 最小値
⇔ $1 , 2 , .... , N$ の集合から区間に含まれる要素を削除して残った最小の値 と解釈できます。
区間はスライドするため、この区間に含まれない要素の集合を管理するためには任意の要素の削除 および 追加 を高速で行えるデータ構造が必要になります。
また、最小値の取得 も高速に行わなければなりません。具体的には、考えなければいけない区間が $N-M+1$ 個あるので、各操作を $O(1)$ 最高でも $O(logN)$ で行う必要があります。
したがって、これらの操作を全て $O(logN)$ で実行可能な 平衡二分探索木 で管理します ! 。。。。と言いたいところですが、残念ながら python の標準ライブラリには平衡二分探索木がありません。代用で SortedSet がありますが、各操作が $O(√N)$ なので厳しいです。
そこで 優先度付きキュー を使います。優先度付きキューからの要素の削除は、その値が最小値でない場合 $O(N)$ となってしまいますが、頑張って工夫することでこのボトルネックを解消することが可能です。
削除の工夫
まず、「 削除を繰り返して残った最小の値 」、これを求める最小値として考えるのではなく、「 現在の最小値を削除すべきであれば削除する、ことを繰り返した現在の最小値 」を求める最小値として考えます。
このようにすれば、優先度付きキューから最小値以外の要素を削除しないで済むので、ボトルネックが解消されることになります。
ただし、これを実現するためには要素の削除を一旦保留しておく必要があるので、この保留中の集合 ：削除待ち集合 を追加で管理します。
以上の操作のイメージはこのようになります。
追加
区間のスライドによって区間からは要素があふれます。もし区間に同じ値の要素がなければ、区間に含まれない要素の集合にこれを追加しなければなりません。
ただし、この時点でまだこの要素が削除待ちである場合、優先度付きキューからこの値は削除されていません。
したがって、追加すべき要素が削除待ち集合になければ追加し、あれば削除待ち集合から削除を行います。
一方、このようにスライドであふれた要素と同じ値の要素が区間に含まれている場合には、区間に含まれない要素は増えません。
この違いを検出するために、各要素の区間内における最も右側の位置を管理します。これによって区間からあふれた際、その要素を追加する必要があるのかを確認することが可能になります。
コード
pypy3
1583 ms/ 2000 ms
from collections import defaultdict
from heapq import heapify,heappop,heappush
N,M=[int(nm) for nm in input().split()]
A=[int(a) for a in input().split()]
# 区間に含まれない要素を 優先度付きキュー で管理
H=list(range(2*10**6))
heapify(H)
# 区間に含まれる mex より大きな要素を管理 (ヒープからの削除待ち要素)
wait_dis=set()
# 各要素の区間内における最も右側の位置を管理。含まれなければ負値
D=defaultdict(lambda : -10)
ans=2*10**6
for i in range(N):
wait_dis.add(A[i])
D[A[i]]=i
if i >= M:
# 区間からあふれる要素と同じ値の要素が区間になければ
if D[A[i-M]] == i-M:
D[A[i-M]]=-10
# 出現しなくなるので、集合から削除 または ヒープに追加
if A[i-M] in wait_dis:
wait_dis.discard(A[i-M])
else:
heappush(H,A[i-M])
if i+1 >= M:
# 現在の最小値を削除すべきなら削除
while H[0] in wait_dis :
wait_dis.discard(H[0])
heappop(H)
ans=min(ans,H[0])
print(ans)
追記 230215
最小値以外を削除しない工夫のために削除待ち集合を用いたが、実は単なる個数管理でも十分であることに気が付いた。
むしろ、削除待ち集合を扱うことで生じていた面倒がなくなって楽。具体的には優先度付きキューに重複を許した追加を行わないために設定していた条件を全て考える必要がなくなる。
削除を行う際に if 区間に含まれる現在の最小値の個数が 0 でない とすることで、優先度付きキューに重複して要素が格納されていても問題なく全て削除できるようになるからである。
コード
pypy3
1587 ms/ 2000ms
from collections import defaultdict
from heapq import heapify,heappop,heappush
N,M=[int(nm) for nm in input().split()]
A=[int(a) for a in input().split()]
# 区間に含まれない要素を 優先度付きキュー で管理
H=list(range(2*10**6))
heapify(H)
# 各要素の区間内における個数
D=defaultdict(int)
ans=2*10**6
for i in range(N):
D[A[i]]+=1
if i >= M:
D[A[i-M]]-=1
heappush(H,A[i-M])
if i+1 >= M:
# 現在の最小値を削除すべきなら削除
while D[H[0]]!=0:
heappop(H)
ans=min(ans,H[0])
print(ans)
別解➀
$mex$ : 区間に含まれない & 最小値
⇔ 区間に含まれる要素を ∞ (最小値に影響させない)とした 1 ～ N までの最小値 と解釈することができます。
したがって、➀ 区間最小値出力 ➁ 区間に含まれる瞬間、あふれる瞬間に要素を更新
この 2つの操作を高速に行えるデータ構造が求められます。区間取得、一点更新 が $O(logN)$ で実行できる Segment tree がピッタリでしょう。
つまり単位元を $∞$ , 更新関数を $min$ として、区間に含まれる要素の位置には単位元を、含まれない要素の位置にはその index 番目を与えて、$N-M+1$ 個の区間それぞれで $mex$ : $min(\ [1,N+1)\ $) を求めていきます。本方針と同様、区間に含まれる最も右側の位置を defaultdict で管理します。
コード
pypy3
1542 ms/ 2000 ms
メイン
# 単位元
def e():
return 2*10**6
# 更新関数
def op(a,b):
return min(a,b)
ここに セグ木 クラスを書く
if __name__ == "__main__":
from collections import defaultdict
N,M=[int(nm) for nm in input().split()]
A=[int(a) for a in input().split()]
ST=SegTree(op, e, N+10 ,[i for i in range(N+10)])
D=defaultdict(lambda : -10)
ans=2*10**6
for i in range(N):
ST.set(A[i],e())
D[A[i]]=i
if i >= M and D[A[i-M]] == i-M:
D[A[i-M]]=-10
ST.set(A[i-M],A[i-M])
if i+1 >= M :
ans=min(ans,ST.prod(0,N+10))
print(ans)
セグメントツリークラス
# 最下層は (0 + 最下層サイズ)-indexed
class SegTree:
def __init__(self, op, e, n, v=None):
self._n=n
# 演算関数
self._op = op
# 単位元
self._e = e
# n ≦ 2^x を満たす最小のx
self._log = len(f"{(self._n)-1:b}")
# 最下層頂点数 = 2^(桁数)
self._size = 1 << self._log
# テーブルサイズ (全頂点数は 2×最下層サイズ であるが、1-indexdedのための調整で +1する)
# 単位元で初期化
self._d = [self._e()] * (2 * self._size)
if v is not None:
# 最下層の初期化
for i in range(self._n):
self._d[self._size + i] = v[i]
# 上の階層の初期化(親の更新)
for i in range(self._size - 1, 0, -1):
self._d[i] = self._op(self._d[i << 1], self._d[i << 1 | 1])
# 一点更新
def set(self, p, x):
p += self._size
self._d[p] = x
while p:
self._d[p >> 1] = self._op(self._d[p], self._d[p ^ 1])
p >>= 1
# 区間取得
# l,r は0-indexed
def prod(self, l, r):
tmp_l, tmp_r = self._e(), self._e()
l += self._size
r += self._size
while l < r:
# 現在頂点を含める必要がない場合はスルーしても OK (親を使えばよいから)
# そうでない場合は現在頂点を区間に含めて、より内側の親へ移動させる
# ➀ 左端 (左の子なら)
if l & 1:
tmp_l = self._op(tmp_l, self._d[l])
# 内側へ移動
l += 1
# ➁ 右端 (右の子なら) (区間の右端が左の子ならば内側に移動するが、これは r が右の子の場合である)
if r & 1:
# r は区間に含まれない
r -= 1
tmp_r = self._op(self._d[r], tmp_r)
# 親へ移動
l >>= 1
r >>= 1
return self._op(tmp_l, tmp_r)
別解➁
$0 ≦ mex ≦ N+1 $ を満たします。もし、$mex = 0 $ であれば、$N-M+1$ の少なくとも 1つの区間では $0$ が含まれないことになります。
これを判定することは実は難しくなく、$0$ が存在する位置の間隔で判定できます。もし、M 以上の間隔があれば、その中に必ず 0 が含まれない区間が存在することになります。
逆に、すべての間隔で M 未満であれば、すべての区間で 0 を含むので $mex ≒ 0 $ です。
よって、次に$\ mex = 1 $ であるか考えます。
本来であれば $N-M+1$ の少なくとも 1つの区間では $1$ が含まれないこと、に加えすべての区間で 0 を含むこと、が条件ですが後者は既に成立が確認できています。
したがって、前者を $1$ が存在する位置の間隔で判定すれば十分です。
これ以降も、同様にして $mex$ が確定するまで判定し続けます。
あらかじめ各数字の位置を求めておけば、計算量は全体で $O($ 間隔の数 $)$ になります。間隔を簡単に求めるために、どの数字も先頭と末尾に含まれる調整を施しても全体で $N+N$ 個の間隔しかないので、十分高速です。
コード
pypy3
1957 ms/ 2000 ms
N,M=[int(nm) for nm in input().split()]
A=[int(a) for a in input().split()]
# 各数字ごとに位置を求める
points=[[0] for _ in range(2*10**6)]
for i in range(N):
points[A[i]].append(i+1)
# 間隔を簡単に求めるための調整
for i in range(2*10**6):
points[i].append(N+1)
for i in range(2*10**6):
L=len(points[i])
for j in range(1,L):
# ひとつでも M 以上の間隔があれば、その中に i が含まれない区間が必ず存在
if points[i][j]-points[i][j-1]+1-2 >= M:
print(i)
exit()
補足
-
本方針の類題
問題ページ
ネタバレ防止のために、abc 280 ～ 256 までの C ～ E のどれかと記載しておきます。 -
別解 ➀ 参考
【Python】ABC194 解説 -
Segment Tree
【Segtree編】AtCoder Library 解読 〜Pythonでの実装まで〜
コードの Segment Tree クラスはこの記事で紹介されているものをほぼ丸パクリしているので、わからないところはこちらで調べてください。めちゃくちゃ丁寧でわかりやすいです。Segment Tree のお勉強(1)
わかりやすい図によって、操作を視覚的に理解することができます。こちらもおすすめです。 -
別解➁ 参考
公式Youtube 動画
終わり
本方針が自分で考えたもの。
別解➀ はセグ木でやりたかったけど思いつかず、調べて真似した。
別解➁ は動画見るまで考えもしなかった、自分の手札にない。
本方針だけ公開でもよかったが、できるだけコードを落すことに意味があると思うのでいっぱいのせた。
質問やご指摘はこちらまで
Twitter : Waaa1471