ABC217 D - Cutting Woods - かんプリンの学習記録

問題概要
解説
- Pythonでの解法
感想

問題概要

長さ ${L}$ の木材がある．以下の ${Q}$ 個のクエリを処理せよ． ${i}$ 番目のクエリは ${(c_i,x_i)}$ で与えられる．

${c_i=1}$ のとき：木材の左端から ${x_i}$ の地点で木材を切る．
${c_i=2}$ のとき：木材の左端から ${x_i}$ の地点を含む木材の長さを出力する．

${1\leq L\leq 10^9\\ 1\leq Q\leq 2×10^5\\ c_i = 1,2\\ 1\leq x_i\leq L-1}$

解説

現在の木材の区間を管理します． ${[0,L)}$ を初期状態とし， ${[l,r)}$ の木材を ${x(l < x < r)}$ の地点で切ると， ${[l,x)}$ と ${[x,r)}$ の木材になります．
C++のstd::setで ${0,L}$ および切れ目を管理することで， ${x}$ の地点を含む木材の値が両端が取得できます．

具体的にはそれぞれのクエリで以下の操作を行います．

クエリ1： ${x_i}$ をsetに入れる．
クエリ2：setの ${x_i}$ 以上の最小の要素 ${A}$ と ${x_i}$ 以下の最大の要素 ${B}$ を取得し， ${A-B}$ を出力する．

Pythonでの解法

C++にはstd::setというデータ構造があり，クエリ1がinsert，クエリ2がlower_boundで実現できるのですが，PythonにはSTLにsetのようなものがない(Pythonよくわかってないのでもしかしたらある？)ので公式解説にあるようにSegment Tree / Fenwick TreeやUnion-Findで解くのですが，これだと座標圧縮やその他setに比べて考えることが増えて，やるだけのC++に比べて時間がかかってしまうので，やっぱりPythonでもsetのようなものが欲しくなります．
結局，以下の操作が高速に行えるようなものが欲しいです．

${x}$ を挿入
${x}$ 以上の最小要素の取得
${x}$ 以下の最大要素の取得

これらの操作が行えるデータ構造としてBinaryTrieがあります．

以下の記事が分かりやすいです．BinaryTrieはXOR関連の操作も可能ですが今回は使いません．

kazuma8128.hatenablog.com

必要なものを実装すると以下の様になります．計算量は ${O(Q\log\max{x})}$ です．

class BinaryTrie:
    def __init__(self, max_query=2*10**5, bitlen=30):
        n = max_query * bitlen
        self.nodes = [-1] * (2 * n)
        self.cnt = [0] * n
        self.id = 0
        self.bitlen = bitlen
 
    # xの挿入
    def insert(self,x):
        pt = 0
        for i in range(self.bitlen-1,-1,-1):
            y = x>>i&1
            if self.nodes[2*pt+y] == -1:
                self.id += 1
                self.nodes[2*pt+y] = self.id
            self.cnt[pt] += 1
            pt = self.nodes[2*pt+y]
        self.cnt[pt] += 1
 
    # 昇順x番目の値
    def kth_elm(self,x):
        pt, ans = 0, 0
        for i in range(self.bitlen-1,-1,-1):
            ans <<= 1
            if self.nodes[2*pt] != -1 and self.cnt[self.nodes[2*pt]] > 0:
                if self.cnt[self.nodes[2*pt]] >= x:
                    pt = self.nodes[2*pt]
                else:
                    x -= self.cnt[self.nodes[2*pt]]
                    pt = self.nodes[2*pt+1]
                    ans += 1
            else:
                pt = self.nodes[2*pt+1]
                ans += 1
        return ans

    # x以上の最小要素が昇順何番目か
    def lower_bound(self,x):
        pt, ans = 0, 1
        for i in range(self.bitlen-1,-1,-1):
            if pt == -1: break
            if x>>i&1 and self.nodes[2*pt] != -1:
                ans += self.cnt[self.nodes[2*pt]]
            pt = self.nodes[2*pt+(x>>i&1)]
        return ans
 
bt = BinaryTrie()
L,Q = map(int,input().split())
bt.insert(0)
bt.insert(L)
for _ in range(Q):
    c,x = map(int,input().split())
    if c == 1:
        bt.insert(x)
    else:
        p = bt.lower_bound(x)
        print(bt.kth_elm(p)-bt.kth_elm(p-1))

さらに機能を追加したものを以下の記事に載せました．(2021/12/22 追記)

kanpurin.hatenablog.com

また，以下のプログラムもなぜか通る．計算量は ${O(Q^2)}$ のはず．arrayのinsertが速いらしい(未検証)．

import bisect
import array
l,q=map(int,input().split())
a=array.array('i',[0,l])
for _ in range(q):
    c,x=map(int,input().split())
    y = bisect.bisect(a,x)
    if c==1:
        a.insert(y,x)
    else:
        print(a[y]-a[y-1])

感想

データ構造ゲーなのでpython初心者にはキツかったかも．
前もって準備できるのは準備すべきなのでC++のSTLの有名なやつやACLにあるやつはpythonで使い方とともにメモしておいた方がいいですね．