2021-11-27

Python で二次元セグ木を実装した話

はじめに

ライブラリ整備シリーズ第一弾
最近の ABC に出たとある問題でお世話になったので実装しました。
下の例では、minimum query の場合に絞って話をしています。

実装概要

乗せる配列を A (4 × 4) とします。
つまり、セグ木(名前を tree とします)は 8 × 8 配列となります。

初期化

① tree[4][4], A[0][0] を基準として、それぞれ対応するところに値を入れます

f:id:Anko_nasubi:20211127161614p:plain

② 各列・各行について、一次元セグ木と同じ要領で最小値を記録していきます(簡単のため、Python では動かない記法を使っています)

f:id:Anko_nasubi:20211127162304p:plain

③ ②で新しく作った各行 (列でも OK) の値を用いて各列の最小値を記録していきます

f:id:Anko_nasubi:20211127163153p:plain

更新

例として、A[1][2] の値を更新します。

① 一次元セグ木と同じ要領で、ビットシフトなどを用いて値がある行、列の対応する場所を更新します(下図の橙色の部分)

f:id:Anko_nasubi:20211127163747p:plain

② ①の更新で更新した行・列を組み合わせて更新します (この場合、tree[5][1], tree[5][3], tree[1][6], tree[2][6] を更新したので、H = 1 or 2 かつ W = 1 or 3 である場所 (下図の青色の部分) を更新します)

f:id:Anko_nasubi:20211127164241p:plain

クエリ処理

一次元セグ木のときの処理を二重 While を用いて組み合わせます。
例えば、1 <= H < 4, 1 <= W < 4 であるような範囲 (下図の橙色) の最小値は、下図の青色で示した部分の最小値を求めることで求めることができます。

f:id:Anko_nasubi:20211127165338p:plain

ソースコード

計算量

セグ木に乗せる配列の大きさを H × W とすると、

初期化 : O(HW)
更新 : O(log H log W)
クエリ : O(log H log W)

となります。

使い方

`init(self, val, segf = min, ide = 10**18)`

val ... セグ木に乗せる配列です。二次元配列でないと動きません。
segf ide ... デフォルトでは minimum query を扱うときのものになっています。必要に応じて指定してください。

`update(self, h, w, x)`

h 列目 w 行目の値を x に更新します。

`query(h1, h2, w1, w2)`

h ∈ [h1, h2), w ∈ [w1, w2) であるような h, w に関して、それらの最小値 (segf = min の場合) を返します。h1 >= h2 または w1 >= w2 の場合 ide を返します。

class seg2d:
    def __init__(self, val, segf = min, ide = 10**18):
        h = len(val)
        w = len(val[0])
        self.segf = segf
        self.ide = ide
        self.h = 1 << (h - 1).bit_length()
        self.w = 1 << (w - 1).bit_length()
        self.tree = [[ide] * (2 * self.w) for _ in range(2 * self.h)]
        for i in range(h):
            for j in range(w):
                self.tree[self.h + i][self.w + j] = val[i][j]
        for i in range(h):
            for j in range(self.w - 1, 0, -1):
                self.tree[self.h + i][j]=self.segf(self.tree[self.h + i][j * 2], self.tree[self.h + i][j * 2 + 1])
        for i in range(w):
            for j in range(self.h - 1, 0, -1):
                self.tree[j][self.w + i] = self.segf(self.tree[j * 2][self.w + i], self.tree[j * 2 + 1][self.w + i])
        for i in range(self.h - 1, 0, -1):
            for j in range(self.w - 1, 0, -1):
                self.tree[i][j] = self.segf(self.tree[i][j * 2], self.tree[i][j * 2 + 1])
    def update(self, h, w, x):
        h += self.h
        w += self.w
        self.tree[h][w] = x
        h2 = h
        while h2 > 1:
            self.tree[h2 >> 1][w] = self.segf(self.tree[h2][w], self.tree[h2 ^ 1][w])
            h2 >>= 1
        w2 = w
        while w2 > 1:
            self.tree[h][w2 >> 1] = self.segf(self.tree[h][w2], self.tree[h][w2 ^ 1])
            w2 >>= 1
        h2 = h
        while h2 > 1:
            w2 = w
            while w2 > 1:
                self.tree[h >> 1][w2 >> 1] = self.segf(self.tree[h >> 1][w2], self.tree[h >> 1][w2 ^ 1])
                w2 >>= 1
            h2 >>= 1
    def query(self, h1, h2, w1, w2):
        ret = self.ide
        h1 += self.h
        h2 += self.h
        w1 += self.w
        w2 += self.w
        while h1 < h2:
            w3 = w1
            w4 = w2
            if h1 & 1:
                while w3 < w4:
                    if w3 & 1:
                        ret = self.segf(ret, self.tree[h1][w3])
                        w3 += 1
                    if w4 & 1:
                        ret = self.segf(ret, self.tree[h1][w4 - 1])
                    w3 >>= 1
                    w4 >>= 1
                h1 += 1
            w3 = w1
            w4 = w2
            if h2 & 1:
                while w3 < w4:
                    if w3 & 1:
                        ret = self.segf(ret, self.tree[h2 - 1][w3])
                        w3 += 1
                    if w4 & 1:
                        ret = self.segf(ret, self.tree[h2 - 1][w4 - 1])
                    w3 >>= 1
                    w4 >>= 1
            h1 >>= 1
            h2 >>= 1
        return ret

2021-09-17

Python で非再帰 DFS を楽に実装したい話

ABC218 で非常にお世話になったので...

実装概要

(返り値・再帰のときの引数にあたるものを記録しておくための配列を作る)
各頂点で処理をするとき、
- はじめに [現在処理をしている頂点, 0] を deque に append する
- その後、移動可能頂点があるならば、 [移動先の頂点, 1] を順番に deque に append していく
(返り値がある場合、que.pop()[1] が 0 ならば移動した頂点から返り値を収集し(配列に記録してある)、処理を行う)

こうすることで、順番に pop するだけで再帰と同じ操作を行うことができる。
言葉だけでは分かり辛いため、以下のような木を探索することを考える。 f:id:Anko_nasubi:20210917133608p:plain 頂点番号が探索順となるように探索してみると、deque の中身は

[[1, 1]]
#頂点1の処理
[[1, 0], [5, 1], [2, 1]]
#頂点2の処理
[[1, 0], [5, 1], [2, 0], [4, 1], [3, 1]]
#頂点3の処理&後処理
[[1, 0], [5, 1], [2, 0], [4, 1]]
#頂点4の処理&後処理
[[1, 0], [5, 1], [2, 0]]
#頂点2の後処理・頂点5の処理
[[1, 0], [5, 0], [6, 1]]
#頂点6の処理&後処理・頂点5の後処理・頂点1の後処理
[]

というようになる。

実装

この実装方法を用いて、試しに ABC165 F - LIS on Tree を解いてみる。

コード例 (PyPy3, 672ms)(提出先)

from collections import deque

n = int(input())
a = list(map(int,input().split()))
path = [[] for _ in range(n)]
for i in range(n-1):
    u, v = map(int,input().split())
    path[u-1].append(v-1)
    path[v-1].append(u-1)

lis = [10**10] * n

def update(x):
    ok = n
    ng = 0
    while abs(ok-ng) > 1:
        ind = (ok + ng) // 2
        if lis[ind-1] >= x:
            ok = ind
        else:
            ng = ind
    return ok

que = deque()
que.append([0, 1])
seen = [False] * n
seen[0] = True
leng = 0
ans = [0] * n
update_memo = [[0] * 3 for _ in range(n)]

while len(que):
    task = que.pop()
    node = task[0]
    if task[1]:
        upd = update(a[node])
        update_memo[node] = [upd-1, lis[upd-1], leng]
        leng = max(leng, upd)
        lis[upd-1] = a[node]
        ans[node] = leng
        que.append([node, 0])
        for i in path[node]:
            if not seen[i]:
                seen[i] = True
                que.append([i, 1])
    else:
        lis[update_memo[node][0]] = update_memo[node][1]
        leng = update_memo[node][2]

for i in ans:
    print(i)