はじめに

以前、遅延セグメント木を作ったときは、コードを書き上げることに注力してしまい、深い理解に至ってなかったので、今回は作った遅延セグメント木をもう少しイジって、汎用性を増すように一般化を進めたい。
前回は、Intのみの狭い世界だったからね。

一般化を進めた遅延セグ木

元ネタは、Intの世界だけで一般化した遅延セグ木であり、区間取得クエリとして「区間最大値クエリ」、区間変更クエリとして「区間加算クエリ」、を行うもの。
通常のセグ木におけるモノイドの型をMonoidとし、遅延セグ木におけるモノイドの型をActMonoidとした。また、左作用するモノイド作用(:ActMonoid,:Monoid)->Monoidをflzとしている。元のflzの引数では、第1引数にMonoidを持ってきてたけど、左作用って事を考慮して、第1引数をActMonoidに入れ替えたよ。
また、一点更新関数のupdateも汎用化する同時に名称をpointUpdateに変え、区間集計クエリ関数queryもareaQueryに名称変更しました。コレでスッキリ！
あと、一点評価関数のpointQueryを新設したよ。これは、クエリーの前後でtreeもlazyも変化しないよ。葉部分[n...]のlazyをtreeに作用させた値を返すだけ。
なお、ActMonoidは作用素モノイドと呼ぶよ。

import Foundation

class LazySegTree {

    // 遅延セグ木設定場所
    typealias Monoid = Int
    typealias ActMonoid = Int 
    private var base: Monoid = -99 // maxクエリなので、負の大きい数
    private var lzBase: ActMonoid = 0 // 加算クエリなので、0
    private var biOpe:(_ l:Monoid,_ r:Monoid)->Monoid = max // maxクエリなので、max
    private var lzOpe:(_ u:ActMonoid,_ d:ActMonoid)->ActMonoid = (+)
    private var flz:(_ a:ActMonoid,_ m:Monoid,_ len:Int)->Monoid = {$0 + $1 + $2 * 0} // ActMonoidとMonoidが同型であるため、「+」が使える

    // 追加固有関数
    var add:(Int,Int,Int)->() { return {self.areaOpe($0,$1,$2)} } // 区間変更クエリが加算なので、add(a,b,x)（範囲a..<bにxを加算）で呼べるようにする。
    var update:(Int,Int)->() {
        return {
            self.areaOpe($0,$0+1,$1 - (self.query($0,$0+1) == self.base ? 0 : self.query($0,$0+1)))// add(i,i+1,x-query(i,i+1))と書くのが面倒なので、update関数を作成
        }
    }

    // 以降は変更不要なはず
    var tree: [Monoid]
    var lazy: [ActMonoid]
    var n: Int
    private func biOpe(_ i:Int)->Monoid{ // 引数1つのときの関数オーバーロード
        return biOpe(tree[2*i],tree[2*i+1])
    }
        
    init(_ ms: [Monoid]) {
        let sz = ms.count
        n = 1
        while n < sz { n *= 2 }
        self.tree = Array(repeating: base, count: 2*n)
        self.lazy = Array(repeating: lzBase, count: 2*n)
        buildTree(ms)
    }

    init(_ cnt: Int, v: Monoid? = nil) {
        let ms = [Monoid](repeating: v ?? self.base, count: cnt)
        n = 1
        while n < cnt { n *= 2 }
        self.tree = Array(repeating: base, count: 2*n)
        self.lazy = Array(repeating: lzBase, count: 2*n)
        buildTree(ms)
    }

    private func buildTree(_ ms: [Monoid]) {
        let sz = ms.count
        for i in 0..<sz {
            tree[n+i] = ms[i]
        }
        for i in stride(from: n-1, through: 1, by: -1) {
            tree[i] = biOpe(i)
        }
    }
    
    private func eval(_ k: Int, _ l: Int, _ r: Int) {
        if lazy[k] != 0 {
            tree[k] = flz( lazy[k] ,tree[k] , r-l )
            if r - l > 1 {
                lazy[2*k] = lzOpe( lazy[k] , lazy[2*k] )
                lazy[2*k+1] = lzOpe( lazy[k] , lazy[2*k+1] )
            }
            lazy[k] = lzBase
        }
    }

    func areaOpe(_ a: Int, _ b: Int, _ x: ActMonoid, _ k: Int = 1, _ l: Int = 0, _ r: Int = -1) {
        var r = r
        if r < 0 { r = n }
        eval(k, l, r)
        if b <= l || r <= a { return }
        if a <= l && r <= b {
            lazy[k] = lzOpe( lzBase , x )
            eval(k, l, r)
        } else {
            areaOpe(a, b, x, 2*k, l, (l+r)/2)
            areaOpe(a, b, x, 2*k+1, (l+r)/2, r)
            tree[k] = biOpe(k)
        }
    }

    func areaQuery(_ a: Int, _ b: Int, _ k: Int = 1, _ l: Int = 0, _ r: Int = -1) -> Monoid {
        var r = r
        if r < 0 { r = n }
        eval(k, l, r)
        if b <= l || r <= a { return base }
        if a <= l && r <= b { return tree[k] }
        let vl = areaQuery(a, b, 2*k, l, (l+r)/2)
        let vr = areaQuery(a, b, 2*k+1, (l+r)/2, r)
        return biOpe(vl,vr)
    }
    
    func pointUpdate(_ i: Int, v: Monoid) { // 一点更新
        var k = i + n
        lazy[k] = lzBase
        tree[k] = v
        while k > 1 {
            k /= 2
            tree[k] = biOpe(k)
        }
    }

    func pointQuery(_ i: Int)-> Monoid { // 一点評価
        var k = i + n
        var ans = tree[k]
        while k > 0 {
            ans = flz( lazy[k] ,ans , 1)        
            k /= 2
        }
        return ans
    }

    /// 下記は、treeやlazyがIntである前提で作ってるけど、とりあえず検証用に放置。モノイドがIntでないときはエラーの原因になるので、捨ててよし。
    /// 捨てるときは、import Foundationも不要になるよ。
    func prtIndex() {
        let nums = (1..<(2*n)).map { String(format: "%3d", $0) }.joined(separator: ",")
        print("(",nums,")")
    }    
    
    func prtTree() {
        let nums = tree[1...].map { String(format: "%3d", $0) }.joined(separator: ",")
        print("[",nums,"]")
        // print(tree[1...])
    }
    func prtLazyTree() {
        let nums = lazy[1...].map { String(format: "%3d", $0) }.joined(separator: ",")
        print("[",nums,"]")
        // print(lazy[1...])
    }

    func prtAry() {
        let nums = tree[n...].map { String(format: "%3d", $0) }.joined(separator: ",")
        print("[",nums,"]")    
        // print(tree[n...])
    }

}

シン遅延セグ木で遊んでみる

せっかく、Int以外のモノイドを使えるようになったので、下記のような遅延セグ木を作ってみる。
- 分析対象の配列
  - 集合：Set<Int>
- 区間取得クエリ
  - 区間のSetに含まれる要素の最大値を取得する。
    - ただし、モノイドとするためには、二項演算(biOpe)の引数も戻り値もモノイド(集合)とする必要があるから、biOpeとしては、集合を統合(union)し、biOpeが終わった後に、統合後の集合からmaxを取ることにしよう。あと、モノイドの単位元は空集合になるね。
- 区間変更クエリ
  - 区間のSetの要素をx:Int倍する
- その他
  - pointUpdateだと、新しい集合を指定することになるけど、これとは別に対象配列のインデックスiにある集合に、値vを追加するins(i,v)を追加。
  - 同様に、iにある集合の最大値を削除するrmMax(i)を追加。
ちなみに、モノイド部分を整理するとこんな感じ。

	型	二項演算子	単位元	外部演算(左作用)
セグ木(モノイド)	Set<Int>	\$0.union($1)	空集合	-
遅延木(作用モノイド)	Int	(*)	1	{mlt,set in Set(set.map{mlt * $0})}

じゃあ、作ってみようか。とはいえ、変更するのは、「遅延セグ木設定場所」「追加固有関数」の箇所だけ。

    // 遅延セグ木設定場所
    typealias Monoid = Set<Int>
    typealias ActMonoid = Int 
    private var base: Monoid = Set<Int>() // unionクエリなので、空集合
    private var lzBase: ActMonoid = 1 // 乗算クエリなので、0
    private var biOpe:(_ l:Monoid,_ r:Monoid)->Monoid = {$0.union($1)} // unionクエリ
    private var lzOpe:(_ u:ActMonoid,_ d:ActMonoid)->ActMonoid = (*)
    private var flz:(_ a:ActMonoid,_ m:Monoid,_ len:Int)->Monoid = {mlt,set,_ in Set(set.map{mlt * $0})} // モノイド作用に相当する

    // 追加固有関数
    func maxq(_ a:Int,_ b:Int)->Int{
        let inf = -99
        return (areaQuery(a,b).max() ?? inf)
    }    
    func mltp(_ a:Int,_ b:Int,_ m:Int){
        areaOpe(a,b,m)
    }
    func ins(_ i:Int,_ v:Int){
        var leaf = areaQuery(i,i+1) // insert前に、遅延評価を済ませておく
        leaf.insert(v)
        pointUpdate(i,leaf)
    }
    func rmMax(_ i:Int){
        var leaf = areaQuery(i,i+1) // 定数倍には、負数も含むから、areaQueryが必要
        guard let mx = leaf.max() else {return}
        leaf.remove(mx)
        pointUpdate(i,leaf)        
    }
    
    func prtTrees() { // 分析用print関数
        for i in 1..<2*n {
            print(i,lazy[i],tree[i])
        }
    }

また、実際に動かしてみたよ。

let sets:[Set<Int>] = [ // 6個の集合
[1,2,3], // 0
[10,20,30,40], // 1
[100,200], // 2
[1000,2000,3000], // 3
[10000], // 4
[1,11,111,1111,11111] // 5
]

print("生成直後")
var tree = LazySegTree(sets)
tree.prtTrees()

print("\n範囲0..<4(ノード8..<12)を8倍した後、インデックス3(ノード11)に5000を追加")
tree.mltp(0,4,8)
tree.ins(3,5000)
tree.prtTrees()

print("\n範囲2..<6の最大値")
print(tree.maxq(2,6))
tree.prtTrees()

print("\n範囲0..<2(ノード8..<10)を-3倍した後、インデックス1(ノード9)の最大値を削除")
tree.mltp(0,2,-3)
tree.rmMax(1)
tree.prtTrees()

各ノードの位置は下図を見てね
出力
- コードを見て貰えば分かるけど、数字部分は、
  - 一番左が、セグ木のインデックス(k)
  - その次が、遅延評価されるlazy[k]
  - 最後が、セグ木のtree[k]

生成直後
1 1 [10, 2000, 20, 10000, 100, 2, 30, 40, 3, 1111, 3000, 1, 200, 1000, 11111, 111, 11]
2 1 [100, 1000, 3000, 2000, 20, 30, 1, 2, 3, 10, 40, 200]
3 1 [10000, 1111, 11111, 1, 11, 111]
4 1 [20, 30, 1, 2, 3, 10, 40]
5 1 [100, 200, 3000, 1000, 2000]
6 1 [10000, 1111, 11111, 1, 11, 111]
7 1 []
8 1 [3, 2, 1]
9 1 [20, 40, 10, 30]
10 1 [100, 200]
11 1 [2000, 3000, 1000]
12 1 [10000]
13 1 [1, 1111, 111, 11, 11111]
14 1 []
15 1 []

範囲0..<4(ノード8..<12)を8倍した後、インデックス3(ノード11)に5000を追加
1 1 [8, 24, 1600, 320, 160, 800, 5000, 10000, 80, 16, 24000, 1111, 8000, 1, 16000, 240, 11111, 111, 11]
2 1 [8, 24, 1600, 320, 160, 800, 5000, 80, 16, 24000, 8000, 16000, 240]
3 1 [1, 111, 11111, 10000, 1111, 11]
4 1 [8, 24, 160, 80, 16, 240, 320]
5 1 [8000, 24000, 16000, 800, 1600, 5000]
6 1 [10000, 1111, 11111, 1, 11, 111]
7 1 []
8 8 [3, 2, 1]
9 8 [20, 40, 10, 30]
10 1 [800, 1600]
11 1 [8000, 5000, 24000, 16000]
12 1 [10000]
13 1 [1, 1111, 111, 11, 11111]
14 1 []
15 1 []

範囲2..<6の最大値
24000
1 1 [240, 24, 80, 1, 11111, 11, 8, 16, 24000, 10000, 8000, 1600, 16000, 160, 320, 800, 5000, 111, 1111]
2 1 [160, 80, 24, 800, 16000, 5000, 320, 8000, 24000, 240, 8, 1600, 16]
3 1 [11111, 10000, 111, 1, 1111, 11]
4 1 [320, 8, 24, 240, 16, 160, 80]
5 1 [1600, 24000, 5000, 8000, 16000, 800]
6 1 [11111, 1, 111, 10000, 1111, 11]
7 1 []
8 8 [3, 2, 1]
9 8 [20, 40, 10, 30]
10 1 [800, 1600]
11 1 [8000, 5000, 24000, 16000]
12 1 [10000]
13 1 [1, 1111, 111, 11, 11111]
14 1 []
15 1 []

範囲0..<2(ノード8..<10)を-3倍した後、インデックス1(ノード9)の最大値を削除
1 1 [-720, -960, 1, 11111, 11, -48, 24000, 10000, 8000, 1600, 16000, -480, -72, 5000, 800, -24, 111, 1111]
2 1 [-72, -480, -720, 8000, 16000, 5000, -960, 1600, 800, -24, -48, 24000]
3 1 [1, 11111, 111, 1111, 11, 10000]
4 1 [-960, -72, -48, -480, -720, -24]
5 1 [1600, 8000, 800, 5000, 24000, 16000]
6 1 [11111, 1, 111, 10000, 1111, 11]
7 1 []
8 1 [-48, -72, -24]
9 1 [-960, -480, -720]
10 1 [800, 1600]
11 1 [8000, 5000, 24000, 16000]
12 1 [10000]
13 1 [1, 1111, 111, 11, 11111]
14 1 []
15 1 []

うむ。想定通りの動きとなってるよ！

おわりに

以前、遅延セグ木の投稿をしたときは、全てIntの世界で発展性が乏しかったから、今回、汎用的なコードに書き換えられて気分がスッキリした。
これでやっと、コンテストで戦える武器になったかな。