@keymoon氏の【全方位木DP】のRuby実装

【全方位木DP】明日使える便利な木構造のアルゴリズム
上記事の Ruby への単なる移植です。

コード

API はブロックを使うなど、多少 Ruby らしくしたつもりです。operateの処理は、Ruby では Array#max が Array のメソッドなので、ちょっと苦肉の策になっています。@operateを attr_writer の対象にしておきましたが、実際は :max を :min にするくらいで、それほどの汎用性を求められないでしょう。

Ruby

class ReRooting
  def initialize(identity, operate = :itself, &operate_node)
    @identity = identity
    @operate = ->(a, b) {[a, b].__send__(operate)}
    @operate_node = operate_node
  end
  attr_writer :operate

  def rerooting(node_count, edges)
    @node_count = node_count

    @adjacents = Array.new(@node_count) {[]}
    @index_for_adjacents = Array.new(@node_count) {[]}

    edges.each do |edge|
      @index_for_adjacents[edge[0]] << @adjacents[edge[1]].size
      @index_for_adjacents[edge[1]] << @adjacents[edge[0]].size
      @adjacents[edge[0]] << edge[1]
      @adjacents[edge[1]] << edge[0]
    end

    @dp = @adjacents.map {|ary| [ary.size]}
    @res = Array.new(@adjacents.size)    #result

    case
    when @node_count > 1
      init
    when @node_count == 1
      @res[0] = @operate_node.(@identity, 0)
    end
    @res
  end

  private

  def init
    parents = Array.new(@node_count)
    order = []

    stack = [0]
    parents[0] = -1

    until stack.empty?
      node = stack.pop
      order << node
      @adjacents[node].each do |adjacent|
        next if adjacent == parents[node]
        stack << adjacent
        parents[adjacent] = node
      end
    end

    #fromLeaf
    (order.size - 1).downto(1) do |i|
      node = order[i]
      parent = parents[node]

      accum = @identity
      parent_index = Object.new
      @adjacents[node].each_index do |j|
        if @adjacents[node][j] == parent
          parent_index = j
        else
          accum = @operate.(accum, @dp[node][j])
        end
      end
      @dp[parent][@index_for_adjacents[node][parent_index]] =
         @operate_node.(accum, node)
    end

    #toLeaf
    order.each do |node|
      accum = @identity
      accums_from_tail = Array.new(@adjacents[node].size)
      accums_from_tail[-1] = @identity

      (accums_from_tail.size - 1).downto(1) do |j|
        accums_from_tail[j - 1] = @operate.(@dp[node][j], accums_from_tail[j])
      end

      accums_from_tail.each_index do |j|
        @dp[@adjacents[node][j]][@index_for_adjacents[node][j]] =
           @operate_node.(@operate.(accum, accums_from_tail[j]), node)
        accum = @operate.(accum, @dp[node][j])
      end
      @res[node] = @operate_node.(accum, node)
    end
  end
end

例えばこんな感じで使います。ノードは 0 から始まる Integer の連番にして下さい。ノードの個数をnode_countで指定します。

sample.rb

edges = [[0, 1], [0, 4], [0, 6], [1, 2], [1, 3], [4, 5],
         [6, 10], [6, 11], [6, 7], [7, 8], [7, 9]]
node_count = 12

res =  ReRooting.new(0, :max) {|value, id| value + 1}
                .rerooting(node_count, edges)
puts res.map(&:pred)

おかしなところがあったら教えて下さい。