セグメント木の実装

前回記事：Juliaでセグメント木（segment-tree）① - 導入

さて、前回はセグメント木の概説をしました。今回からいよいよ実際のJuliaのコードを書いていきたいと思います。少しだけおさらいしましょう。セグメント木の実装には最低限、初期化、区間取得、更新の三つの部品が必要となります。今回はさらにオプションとして一点取得（数列のある項の現在の値を取得）を$O(1)$で行う関数を加えて実装していきたいと思います。

初期化 : mutable struct SegmentTreeのコンストラクタとして実装。
区間取得 : get関数として実装。
更新 : update関数として実装。
一点取得（オプション） : 必須ではないが、頻用するので、特別に$O(1)$で取得できるように実装する。

今回は初期化に絞って話をしたいと思います。

初期化

初期化の方法は前回書いた通りです。これをJuliaのコードで実装します。

mutable struct

structはある用途で使う変数を一つの単位にまとめたもので、CやC++にもあり、特に難しい概念ではないと思います。セグメント木に使うデータを入れる変数をまとめたstructを作っていこうと思います。しかし、注意点としてJuliaにおいてはただstruct hogeと書くと、非可変型のstructを定義することになります。すなわち、一度コンストラクタでstructのメンバ変数に値を入れると変更が不可能になります。セグメント木は動的に各ノードの値を更新するので、変更不可能だと困ります。そこで、mutable修飾子をつけることで可変型のstructを作ることができます。

mutable struct SegmentTree

end

※なぜかJulia のシンタックスハイライトでmutable struct hogeと打つと崩れてしまうのでハイライトなしでご容赦ください

空のSegmentTreeのmutable structが作成されました。では、ここでメンバとなる変数を付け足していこうと思います。

必要なメンバ変数は

モノイドの単位元
モノイドの演算
各ノードの値を格納する配列
元の数列の長さ
元の数列の長さ以上の整数で$2^s$の形で表せる最小の整数
上記の指数s

の６つです。足し算だけでなく、任意のモノイドに対応できるような汎用性の高いセグメント木を作りたいので、メンバ変数にどのようなモノイドか、という情報も格納します。この宣言はmutable structスコープ内に書けば大丈夫です。

mutable struct SegmentTree
    op::Function        #operator : 演算子
    ide::Int64          #単位元
    node::Array{Int64}  #nodeを格納するArray
    sz::UInt32          #元の数列の長さ
    n::UInt32           #元の数列の長さ以上の最小の2の累乗数
    s::UInt8            #2^s = n
end

Juliaではhoge::fugaと書いて、hogeの型はfugaですよ、と宣言します。ちなみにJuliaは型推論を行ってくれるので、厳密に一つ一つ型を宣言する必要はないはずです。ただ今回は競技プログラミングで使おうという目的もあり、コンパイル時間が短くて済む（らしい）型指定を比較的きっちりやっています（これがbetter Juliaかと聞かれると自信がないですが）。

opはoperator、すなわち、演算です。二項演算は要するに二つの要素を引数に取る関数と考えることができます。Juliaでは関数も変数のように扱え、関数はFunctionという抽象型から派生した具象型を持ちます（関数ごとに異なる具象型がいちいち定義される？）。というわけでFunction型で宣言しておけば、その子の型を持つであろう関数を受け取れるわけですね。

他は整数型にしていますが、それぞれがとりうる値を元に使用するメモリ量を変えています。指数sなんかが1000や10000の値を取るとは到底思えない（そんな値が来たらnode分のメモリ確保がまずできません）のでUInt8で宣言しています。

パラメトリック型

また、Juliaの機能の一つとして、便利なのがパラメトリック型です。セグメント木で扱いたい対象とする数列は整数のものだけではありません。浮動小数点数なども和や積に関してモノイドですし、演算子と単位元が定義できるのであれば一般的に数と言われるものじゃなくてもいいわけです。そこで、セグメント木を使うときその都度扱うデータの型を変更できるようにする仕組みとしてパラメトリック型を使います。

mutable struct SegmentTree{T}
    op::Function        #operator : 演算子
    ide::T              #単位元
    node::Array{T}      #nodeを格納するArray
    sz::UInt32          #元の数列の長さ
    n::UInt32           #元の数列の長さ以上の最小の2の累乗数
    s::UInt8            #2^s = n
end

Int64をTに変えました。C++でいうテンプレートですね。SegmentTreeを初期化する際に型まで自在に変えられるようになりました。

コンストラクタ

では、いよいよコンストラクタを通じて、セグメント木の初期化を実装します。コンストラクタを入れたこのコードがmutable struct SegmentTreeの完成版となります。

mutable struct SegmentTree{T}
    op::Function        #operator : 演算子
    ide::T              #単位元
    node::Array{T}      #nodeを格納するArray
    sz::UInt32          #元の数列の長さ
    n::UInt32           #元の数列の長さ以上の最小の2の累乗数
    s::UInt8            #2^s = n

    #constructor
    function SegmentTree{T}(array::Array{T},op::Function,ide::T) where T
        sz::UInt32 = length(array)   #数列の長さを格納
        n::UInt32 = 1                #nの初期化
        s::UInt8 = 1                 #sの初期化
        while n < sz                 #nがszを超えるまでループ
            n *= 2
            s += 1
        end
        node::Array{T} = [ide for i=1:(2*n-1)]     #nodeの初期化。単位元にすることに注意。 

        for i=1:sz                                 #最下段に与えられた数列を入れる
            node[i+n-1] = array[i]
        end

        for i in n-1:-1:1                          #親ノードに子ノードの和を入れる
            node[i] = op(node[i*2],node[i*2+1])
        end
        new{T}(op,ide,node,sz,n,s)                 #mutable struct SegmentTreeを構築し、戻り値として返す
    end
end

array引数に対象とする数列情報を入れます。op引数にモノイドの演算処理を行う関数、ide引数にモノイドの単位元をいれます。細かい処理を順を追って見て行きます。

セグメント木の大きさの決定

        sz::UInt32 = length(array)   #数列の長さを格納
        n::UInt32 = 1                #nの初期化
        s::UInt8 = 1                 #sの初期化
        while n < sz                 #nがsz以上になるまでループ
            n *= 2
            s += 1
        end

数列の長さ以上の最小の2の累乗数がセグメント木の最下段の長さになります。どういうことか下図に示しましょう。

完全二分木にするために単位元で数列を補完して数列長を2の累乗数にするということが目的です。こうすると、セグメント木をいつでも完全二分木とすることができ、実装がシンプルで楽になります（実は完全二分木化しなくてもできるのですが、そうするとノードの親と子の対応関係の管理等いろいろ面倒なことがでてきます、メモリ上の無駄はたかだか2倍までで抑えられるので完全二分木化するのがベターだと思います）。単位元をいれているので計算結果は変わりません。これでセグメント木の大きさが決まりました。

セグメント木の構築（下から上へ）

        node::Array{T} = [ide for i=1:(2*n-1)]     #nodeの初期化。単位元にすることに注意。 

        for i=1:sz                                 #最下段に与えられた数列を入れる
            node[i+n-1] = array[i]
        end

        for i in n-1:-1:1                          #親ノードに子ノードの演算結果を入れる
            node[i] = op(node[i*2],node[i*2+1])
        end

セグメント木は木ですが、木の構造自体は静的な完全二分木であるため、辺の情報をいちいち変数に格納する必要がありません。$v_i (i=2,3,...,2n-1)$の親は$v_{\lfloor i/2 \rfloor}$という決まりを作った配列を用意してあげれば大丈夫です（Juliaは1-indexedであることに注意！）。完全二分木にしたことで簡単な規則で子と親の対応関係が分かるのは実装上大きなメリットですね。配列の長さが2n-1なのはセグメント木のノード数が$\sum_{k=0}^{s} 2^k = 2^{s+1}-1 = 2 \cdot 2^s -1 = 2n-1$となるためです。具体的には下図のようにindexを振ったノードを用意し、最下段に数列を入れます。

そしてこの規則に従い、親ノードに子ノードの演算（たとえば足し算など）を入れていきます。

はしょっていますが、要するに後ろの方のノードから、既に計算済みの子のノードを利用して自分のノードの値を計算しています。和であれば1+2のように真ん中に演算子を置くのが普通の書き方ですが、どんなモノイドの演算にも対応できるように、op(a::T,b::T)のような、二つの引数を取りT型の数を戻り値として返す関数を与えてあげて、演算規則とします。opが足し算を表すのならop(1,2)は3を返すわけですね。これによりセグメント木の初期化が終わりました。

デフォルトコンストラクタにデータを入れて返す

        new{T}(op,ide,node,sz,n,s)                 #mutable struct SegmentTreeを構築し、戻り値として返す

最後に構築したセグメント木の情報をmutable structのデフォルトコンストラクタに渡してあげます。デフォルトコンストラクタはJuliaでstructを宣言した際に自動的に生成されるコンストラクタです。メンバ変数として宣言された順番通りにnew関数に渡すことで呼び出せます。また、Juliaでは関数の最後に書いた式の値をその関数の戻り値として返す（returnを省略できる）ので、この文を最後に書くことで、new関数の戻り値（すなわちmutable struct SegmentTreeのインスタンス）を返すことができます。

今回は初期化について見て行きましたが、次回は区間取得について見て行こうと思います。誤り等があれば教えていただけると幸いです。

Juliaでセグメント木（segment-tree）② - 初期化