Rubyで書く遅延セグ木

自分用にRubyで非再帰版の遅延セグメント木を書いたので、そのメモと使用時の注意点です。

遅延評価セグ木や非再帰実装についての参考

• AtCoder LibraryのLazy Segtreeの使い方
• 非再帰版の遅延評価セグメント木の実装メモ
• 遅延評価セグメント木をソラで書きたいあなたに

方針

ACLの lazy_segtree に必要な構造/関数 S, F, mapping, composition, op を一つのクラスにまとめ、破壊的変更メソッドとして、実装することでオブジェクト生成を抑えるというのが方針です。

ACLの遅延セグ木をもとに実装を行うとどうしてもオブジェクト生成数が多くなってRubyだとキツかったのでこの形にしました。

具体的には先に上げた必要なデータ、関数をインスタンス変数およびメソッドとして一つにまとめたクラスをNodeとし、以下のように構成します。

S : データの型 Nodeクラスのインスタンス変数を使用して、任意に設定する
F : 操作の型 インスタンス変数 @act のみを用いて設定する。
S mapping(F f, S x): @act を自身に破壊的に反映するメソッドとして#map!を実装します。map! 後 @act は id相当の値にします。

破壊的に反映しないものは@actを反映した値を返すメソッドを実装します。

F composition(F f, F g) : @act と合成するメソッドとして #composite!(f) を実装します。f が後に適用される操作です。アフィン変換など操作の合成が可換でない場合に気を付ける必要があります。
また単純な値でない場合は、オブジェクト生成数を抑えるため、自身を破壊的に変更するよう設定するのが望ましいです。
S op(S a, S b): Node同士をマージして、自身に破壊的に反映するメソッドとして、#merge!(lnode, rnode)を実装します

使い方

AtCoder Library Practice Contest : L - Lazy Segment TreeでACしたコード をもとに説明します。問題/方針についてはAtCoder Library Practice Contest L - Lazy Segment Treeを参考にします。

require "lazy_segment_tree"
class Node
  # require methods
  def initialize; set!(0, 0, 0, 0); end

  attr_reader :act
  def composite!(f); @act ^= f; end
  def map!
    # @actを反映した値にインスタンス変数を変更する
    # @actはid相当の値にする
    # @c0=...としてしまうと、@act=1のとき、c1の値が変わってしまう
    set!(c0, c1, inv, 0)
  end
  def merge!(lnode, rnode)
    # @c0 = ... とすると内部の実装都合でprod時に値が壊れる
    x0 = lnode.c0 + rnode.c0 
    x1 = lnode.c1 + rnode.c1
    x_inv = lnode.inv + rnode.inv + lnode.c1 * rnode.c0
    set!(x0, x1, x_inv, @act)
  end
  
  # safe mapping methods.
  def c0;  @act.zero? ? @c0 : @c1; end
  def c1;  @act.zero? ? @c1 : @c0; end
  def inv; @act.zero? ? @inv : @c1 * @c0 - @inv; end

  # utility methods
  def from(x01);set!(1 - x01, x01, 0, 0); end
  def set!(x0, x1, x_inv, x_act)
    @c0 = x0; @c1 = x1; @inv = x_inv; @act = x_act
  end
end

N, Q = gets.split.map(&:to_i)
A = gets.split.map(&:to_i)
lz = LazySegmentTree.new(Node, N){|i, node| node.from(A[i]) }
ans = []

Q.times do
  t, l, r = gets.split.map(&:to_i)
  # 1-indexed 閉区間を0-indexed 半開区間にあわせる
  l -= 1
  if t == 1
    lz.apply(l, r, 1)
  else
    ans << lz.prod(l, r).inv
  end
end
puts ans

インスタンス変数については、以下のとおりです。

• S
  • @c0 : 0の個数 初期値: 0
  • @c1 : 1の個数 初期値: 0
  • @inv: 転倒数 初期値: 0
• F
  • @act: この区間の01を反転するか? (0: しない, 1: する) 初期値: 0

まず set!, from について説明します。
set! はインスタンス変数をまとめて設定するメソッドです。map!、merge!でも使っていますが、本質的な操作が破壊的な操作ばかりなので、値を控えておくために変数をたくさん用意するより便利な場合が多いです。

from は入力を反映するメソッドです。LazySegmentTreeのコンストラクタをブロック付きで呼び出すとインデクスとノードと一緒に呼び出されますので入力を反映するのにつかいます。

from,set!はインターフェイス的にはなくとも良いのですがあると便利です。

#c0, #c1, #inv は@actを反映したときの値を返します。@actの状況に応じて求める値を返します。
通常、attr_reader でインスタンス変数を直接返して、別の名前をつける方法も考えられますが、個人的にはこれがわかりやすいと思っています。これらもインターフェースとしては必須ではありません。
対して#actは必須です。ただし、通常これはattr_readerで構いません。

map!, composite!(f), merge!(lnode, rnode) コメントに書いてありますが、破壊的なメソッドなので注意が必要です。

例えば map! を以下のように書くと値がおかしくなります。

def map!
  @c0 = c0 # @actが1のとき、#c0は@c1の値を返す
  @c1 = c1 # @actが1なので、#c1は@c0の値を返すが、前行で書き換わっているので
  ...
end

そのため一度、変数にいれて、入れ直す必要があります。#set!(x0,x1,x_inv,x)の便利さがわかるところです。

def map!
  x0 = c0 # @actが1のとき、#c0は@c1の値を返す
  x1 = c1 # @actが1なので、#c1は@c0の値を返すが、前行で書き換わっているので
  ...
  @c0 = x0; @c1 = x1;
end

#merge!(lnode, rnode)も同様です。特に#prodの際には 自身を引数として #merge! に渡している箇所があるため、より複雑ですが、やはり#set!でまとめて渡せるとそういうトラブルはかなり解消できます。

#composite!(f)は操作の合成です。今回はxorです。特に注意点はありません。
可換のため、後述する commutative を true に指定することが可能です。

本体

利用できるのは、以下のみです。

• LazySegmentTree#new(node_class, n){|i, node| ... }
  • コンストラクタ。Nodeクラスと要素数をあたえます
  • ブロックを渡すと要素数分インデクス付きで呼び出します
• #prod(l, r)
  • 0-indexed 半開区間 [l, r) の要素をマージしたnodeを返します
• #apply(l, r, x)
  • 0-indexed 半開区間 [l, r) に操作 x を行います
• #commutative=
  • 操作の合成が可換の時に真を指定します。apply時の@actの伝搬を行わないため、少し高速になります(2sの制限時間に2200ms未満でTLEしているときには期待が持てます)。ただし既定ではfalseです

class LazySegmentTree

  # LazySegmentTree.new(Node, n){|i, node| ... }
  def initialize(node_class, n, &block)
    @commutative = false
    @node_class = node_class
    verify_node_class!
    @size = n
    @offset = (1 << (n - 1).bit_length)
    @node = Array.new(@offset << 1){ node_class.new }
    n.times{|i| block[i, @node[@offset + i]] } if block
    k = @offset
    merge(k) while (k -= 1) > 0
  end
    
  # When @commutative is true, #apply (l,r,x) does not propagate @act values.
  # Advance when composition(f, g) is commutative.
  # default: false
  attr_accessor :commutative
  
  def prod(l, r)
    l += @offset
    r += @offset
    g = adaptive_group(l, r)

    lx = @node_class.new
    rx = @node_class.new
    g.reverse_each{|k| propagate(k) }
    while l < r
      (lx.merge!(lx, @node[l]); l += 1) if l.odd?
      (r -= 1; rx.merge!(@node[r], rx)) if r.odd?
      l >>= 1; r >>= 1
    end
    lx.merge!(lx, rx)
    lx
  end

  def apply(l, r, x)
    l += @offset
    r += @offset
    g = adaptive_group(l, r)

    g.reverse_each{|k| propagate(k) } unless @commutative
    while l < r
      (propagate(l, x); l += 1) if l.odd?
      (r -= 1; propagate(r, x)) if r.odd?
      l >>= 1; r >>= 1
    end
    g.each{|k| merge(k) }
  end

  private
  def verify_node_class!
    require_methods = [:merge!, :composite!, :map!, :act]
    require_methods.reject!{|sym| @node_class.method_defined?(sym) }
    raise "#{require_methods} was not defined." if require_methods.size > 0
  end
  
  def propagate(k, x = nil)
    current = @node[k]
    current.composite!(x) if x
    f = current.act
    if k < @offset
      @node[2 * k].composite!(f)
      @node[2 * k + 1].composite!(f)
    end
    current.map!
  end
  
  def merge(k)
    @node[k].merge!(@node[2 * k], @node[2 * k + 1])
  end

  def adaptive_group(l, r)
    a = l / (l & -l)
    b = r / (r & -r)
    g = []
    while l > 0 && l < r
      g << r if r < b
      g << l if l < a
      l >>= 1; r >>= 1
    end
    (g << l; l >>= 1) while l > 0
    g
  end
end

作ってみた感じ

まあまあ早い/使いやすいけど、コンテストでなんでも通るわけではないので、実際のところは微妙という感じです。
いくつか過去問やりましたが、結構TLEします。

遅延セグ木はrubyではあまり書く人みないなという印象があります。
結構うまく書けてもとTLEする問題が多いし、ぎりぎりなら別の言語で書くほうがよい。
実際、ABCならF問題あたりからなので、別の言語を試し始めている頃なのかなと思います。
そういった状況のなかで方針を試せるくらいのものを整備できたのかなとおもっています。