Pythonでお手軽MultiSet！

はじめに

python勢が一度はぶつかるmultiset、、、

もう何も怖くない！！

背景

実装

BinaryIndexTree（BIT）を基に実装。

• self.List
  • Setに格納される予定のある、重複のない値のリスト
• self.size
  • self.Listの長さ
• self.length
  • multisetの要素数
• self.tree
  • BIT用のリスト、値の個数が入る
• self.kvlist
  • 値をself.ListのIndexに変換する辞書（座圧とも呼ばれる）、1-Indexed

※Setに入る値の種類が分かっている前提でのみ使用可能。

class MultiSet:
    def __init__(self, List):
        self.List   = sorted(set(List))
        self.size   = len(self.List)
        self.length = 0
        self.tree   = [0] * (self.size + 1)
        self.kvlist = {j:i+1 for i,j in enumerate(self.List)}
 
    def add(self, i, x):
        i = self.kvlist[i]
        self.length += x
        while i <= self.size:
            self.tree[i] += x
            i += i & -i
            
    def bisect_left(self, w):
        if w <= 0:return 0
        x=0;r=1;n=self.size
        while r<n: r = r<<1
        while r>0:
            if x+r < n and self.tree[x+r] < w:
                w-=self.tree[x+r]
                x+=r
            r>>=1
        return x
    
    def __getitem__(self, w):
        return self.List[self.bisect_left(w)]
 
    def __len__(self):
        return self.length

使用方法

リストの長さをNとしたときO(log N)で値の追加、削除、参照が可能。

List = [7,3,3,5,5,5,1]
ms = MultiSet(List)

for v in List:
    # vを1個追加
    ms.add(v,1)

print([ms[i]for i in range(1, ms.size+1)]) # [1, 3, 3, 5, 5, 5, 7]

for v in List:
    # vを1個削除
    ms.add(v,-1)

print([ms[i]for i in range(1, ms.size+1)]) # []

応用編

二分探索したいとき

from bisect import bisect_left

List = [7,3,3,5,5,5,1]
ms = MultiSet(List)

for v in List:
    # vを1個追加
    ms.add(v,1)

print(bisect_left(ms,5)) # 4 (1-indexed)

簡単！

検証

ABC241 D問題

1165ms (2s制限)

Q = int(input())
List = []
QueryList = []
for _ in range(Q):
    QueryList += list(map(int,input().split())),
    if QueryList[-1][0] == 1:
      List += QueryList[-1][1],
 
from bisect import bisect_right, bisect_left
 
# 番兵
ms = MultiSet(List + [float('inf')])
ms.add(float('inf'),1)
 
for i,*q in QueryList:
  if i==1:
    x, = q
    ms.add(x,1)
  elif i==2:
    if len(ms)==1:print(-1);continue
    x,k = q
    idx = bisect_right(ms,x) - 1
    print([-1,ms[idx-k+1]][0 < idx-k+1 < len(ms)])
  else:
    if len(ms)==1:print(-1);continue
    x,k = q
    idx = bisect_left(ms,x)
    idx += idx==0
    print([-1,ms[idx+k-1]][0 < idx+k-1 < len(ms)])

0 < hogehoge < len(ms)は1-indexedなのと番兵除外のため

参考記事

BITの実装

BITでの二分探索