PythonでC++のMultisetっぽいことをする(2本のSegment Treeを用いた実装)

はじめに

PythonでC++のMultiset（多重集合）っぽいデータ構造を実装するシリーズ第2弾です。
前回のheapqを用いた実装では多重集合で

・値の挿入
・最大値の取り出し
・最小値の取り出し
・ある値xが含まれているか

といったクエリをそれぞれO(logN)で処理できるものを実装しました。(Nは多重集合のサイズ)

この記事ではこれらのクエリに加え二分探索のようなことが行えるもの、つまり

・ある値x未満で最も大きい数は何か？
・ある値x以上で最も小さい数は何か？

というクエリにもO(logN)で答えられるようなデータ構造を2本のSegment Treeを用いて実装していきたいと思います。

背景

実は先日のABC217のD問題でこの二分探索を行えるデータ構造を必要とする問題が出題され話題になりました。
C++で参加していた人からすると標準ライブラリのsetを使うだけで実装が楽だったかと思いますが
Pythonを使っていた人（特に茶～緑の人）には水色以上のdiffに感じられたかもしれません。

Segment Treeとは

簡単にいうと区間Max,区間Min,区間和 etc.を高速に計算できる便利なデータ構造です。
競技プログラミングでは頻出なので知らない人は解説サイトを参照してください。

コード

長いですがとりあえずコードです。

#疑似multiset:segtreeを用いた実装

#####segfunc#####
def segfunc_min(x, y):
    return min(x, y)        #最小値
def segfunc_max(x, y):
    return max(x, y)        #最大値

#####単位元(ide_ele)#####
ide_ele_min = float('inf')      #最小値
ide_ele_max = -float('inf')     #最大値

class SegTree:
    """
    参考：https://qiita.com/takayg1/items/c811bd07c21923d7ec69
    init(init_val, segfunc, ide_ele): 配列init_valで初期化 O(N)
    update(k, x): k番目の値をxに更新 O(N)
    query(l, r): 区間[l, r)をsegfuncしたものを返す O(logN)
    j番目の要素を取り出すとき->sgt.tree[sgt.num+j]
    """
    def __init__(self, init_val, segfunc, ide_ele):
        """
        init_val: 配列の初期値
        segfunc: 区間にしたい操作
        ide_ele: 単位元
        n: 要素数
        num: n以上の最小の2のべき乗
        tree: セグメント木(1-index)
        """
        n = len(init_val)
        self.segfunc = segfunc
        self.ide_ele = ide_ele
        self.num = 1 << (n - 1).bit_length()
        self.tree = [ide_ele] * 2 * self.num
        # 配列の値を葉にセット
        for i in range(n):
            self.tree[self.num + i] = init_val[i]
        # 構築していく
        for i in range(self.num - 1, 0, -1):
            self.tree[i] = self.segfunc(self.tree[2 * i], self.tree[2 * i + 1])

    def update(self, k, x):
        """
        k番目の値をxに更新
        k: index(0-index)
        x: update value
        """
        k += self.num
        self.tree[k] = x
        while k > 1:
            self.tree[k >> 1] = self.segfunc(self.tree[k], self.tree[k ^ 1])
            k >>= 1

    def query(self, l, r):
        """
        [l, r)のsegfuncしたものを得る
        l: index(0-index)
        r: index(0-index)
        """
        res = self.ide_ele

        l += self.num
        r += self.num
        while l < r:
            if l & 1:
                res = self.segfunc(res, self.tree[l])
                l += 1
            if r & 1:
                res = self.segfunc(res, self.tree[r - 1])
            l >>= 1
            r >>= 1
        return res

class QuasiMultiset:   
    def __init__(self,N):
        self.N = N
        self.cnt = [0]*self.N
        self.sgt_min = SegTree([N-1]*N,segfunc_min,ide_ele_min)
        self.sgt_max = SegTree([0]*N,segfunc_max,ide_ele_max)

    def add(self,x):
        if self.cnt[x]==0:        
            self.sgt_min.update(x,x)
            self.sgt_max.update(x,x)
        self.cnt[x] += 1

    def delete(self,x):
        '''
        値がxの要素を一つ削除
        '''
        self.cnt[x] -= 1
        if self.cnt[x] == 0:
            self.sgt_min.update(x,self.N-1)
            self.sgt_max.update(x,0)


    def popMin(self):
        mi = self.sgt_min.query(0,self.N)
        self.delete(mi)
        return mi

    def popMax(self):
        ma = self.sgt_max.query(0,self.N)
        self.delete(ma)
        return ma

    def searchMin(self):
        return self.sgt_min.query(0,self.N)

    def searchMax(self):
        return self.sgt_max.query(0,self.N)

    def lower_bound(self,x):
        '''
        x未満で最大の数を返す
        '''
        return self.sgt_max.query(0,x)

    def upper_bound(self,x):
        '''
        x以上で最小の数を返す
        '''
        return self.sgt_min.query(x,self.N)

    def contain(self,x):
        '''
        xが存在するか->bool
        '''
        return bool(self.cnt[x])

メゾットの説明

0. インスタンスの作成

インスタンスを作成します。
Nはデータのサイズです。セグメント木を使うのでデータサイズを予め決めなければいけないです。
つまり、multisetに出し入れできる値は0以上N未満に限られます。

multiset = QuasiMultiset(N)

コンストラクタでは
・要素の個数管理用の配列cnt
・区間Minクエリ、区間Maxクエリを処理できるセグメント木1本ずつ
を用意します。セグメント木の初期値は全ての要素が区間MinではN-1、区間Maxでは0です。

def __init__(self,N):
    self.N = N
    self.cnt = [0]*self.N
    self.sgt_min = SegTree([N-1]*N,segfunc_min,ide_ele_min)
    self.sgt_max = SegTree([0]*N,segfunc_max,ide_ele_max)

1. addクエリ

多重集合に値を追加します。

multiset.add(x)

具体的には両方のセグメント木のx番目の値をxに変更し、xのカウンタをインクリメントします。

def add(self,x):
    if self.cnt[x]==0:        
        self.sgt_min.update(x,x)
        self.sgt_max.update(x,x)
    self.cnt[x] += 1

2. deleteクエリ

多重集合から要素xを一つ削除します。

multiset.delete(x)

この操作で多重集合内のxの数が0個になると両方のセグメント木のx番目の値を初期化します。

def delete(self,x):
    self.cnt[x] -= 1
    if self.cnt[x] == 0:
        self.sgt_min.update(x,N-1)
        self.sgt_max.update(x,0)

3. popMinクエリ

多重集合内の要素のうち最小の要素をpopします。

mi = multiset.popMin()

具体的には全体の最小値を取り、その値をdeleteします。

def popMin(self):
    mi = self.sgt_min.query(0,self.N)
    self.delete(mi)
    return mi

4. popMaxクエリ

多重集合内の要素のうち最大の要素をpopします。

ma = multiset.popMax()

具体的には全体の最大値を取り、その値をdeleteします。

def popMax(self):
    ma = self.sgt_max.query(0,self.N)
    self.delete(ma)
    return ma

5. searchMinクエリ

popMinと違って最小値を返すだけです。

mi = multiset.searchMin()

6. searchMaxクエリ

popMaxと違って最大値を返すだけです。

ma = multiset.searchMax()

7. lower_boundクエリ

ある値x 未満で多重集合内に存在する最大の要素を返します。

lb = multiset.lower_bound(x)

具体的には[0,x)の区間Maxを取ります。

def lower_bound(self,x):
    return self.sgt_max.query(0,x)

8. upper_boundクエリ

ある値x 以上で多重集合内に存在する最小の要素を返します。

ub = multiset.upper_bound(x)

具体的には[x,N)の区間Minを取ります。

def upper_bound(self,x):
    return self.sgt_min.query(x,self.N)

9. containクエリ

多重集合内に要素xが存在するかというクエリに対してbool型で返します。

b = multiset.contain(x)

デメリット

今回実装した疑似multisetは前回heapqを用いて実装したものに比べ
二分探索ができるようになった分、上位互換に見えるかもしれませんがデメリットもいくつか存在します。

データサイズが限られている

実装にセグメント木を用いるため最初にデータサイズを指定しないといけないというデメリットが生じます。
しかし、昨今のAtCoderで出される問題はクエリ先読みができる形式のものが多く
座標圧縮を用いることでうまく対処できる可能性が高いです。

インスタンスの作れる数が限られている

これも先ほどと同様にセグメント木の性質から生じるデメリットです。
データサイズNのインスタンスを作成する場合、セグメント木も同時に作成するのでO(N)メモリが必要となります。
その点heapqを用いた実装では空のheap木を作成するだけでいいのでメモリの心配はしなくていいです。

使用例

・ABC217 D. Cutting Woods : 先日のABCで出題され、話題になった問題です。座標圧縮を用いることでクエリ数以下のデータサイズに収められます。あとはクエリを順番に消化、二分探索でxの左右の値を求めるだけです。

・ABC218 H. Red and Blue Lamps : 公式解説、解法1の貪欲解法で使用します。公式解説には書かれていませんが、heapqで取ってきた位置の右隣と左隣を探すときにset上の二分探索が必要になります。