クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「3-3 セグメント木によるRMQの実装」(p.155)

この記事は「プログラミングコンテストチャレンジブック第2版」(蟻本)の
「3-3 セグメント木によるRMQの実装」(p.155)
のコードをPythonで書き直したものとなっています。

コード

import sys
sys.setrecursionlimit(4100000)

MAX_N = 1 << 17
INT_MAX = (1 << 31) - 1

# セグメント木を持つグローバル配列
n = int(input())    # セグ木のサイズ
dat = [0] * (2 * MAX_N - 1)

# 初期化
def init(n_):
    # 簡単のため、要素数を2のべき乗に
    n = 1
    while n < n_:
        n *= 2
    
    # すべての値をINT_MAXに
    for i in range(2 * n - 1):
        dat[i] = INT_MAX
    return n
n = init(n)

# k番目の値(0-indexed)をaに変更
def update(k, a):
    # 葉の節点
    k += n - 1
    dat[k] = a
    # 登りながら更新
    while k > 0:
        k = (k - 1) // 2
        dat[k] = min(dat[k * 2 + 1], dat[k * 2 + 2])

# [a, b)の最小値を求める
# 後ろのほうの引数は、計算の簡単のための引数。
# kは節点の番号、l, rはその節点が[l, r)に対応づいていることを表す。
# したがって、外からはquery(a, b, 0, 0, n)として呼ぶ。
def query(a, b, k, l, r):
    # [a, b)と[l, r)が交差しなければ、INT_MAX
    if r <= a or b <= l:
        return INT_MAX
    
    # [a, b)と[l, r)を完全に含んでいれば、この節点の値
    if a <= l and r <= b:
        return dat[k]
    else:
        # そうでなければ、2つの子の最小値
        vl = query(a, b, k * 2 + 1, l, (l + r) // 2)
        vr = query(a, b, k * 2 + 2, (l + r) // 2, r)
        return min(vl, vr)