# -*- coding: utf-8 -*-

class PaigeTarjan
  Elem = Struct.new(:data, :edges,
                    :block, :prev_elem, :next_elem,
                    :count, :tmp_next)
  class Elem
    def initialize(data)
      super(data, [], nil, nil, nil, nil, nil)
    end
  end
  
  Edge = Struct.new(:src, :count)
  
  Count = Struct.new(:value)
  
  Block = Struct.new(:size, :next_elem,
                     :ublock, :prev_block, :next_block, :assoc)
  class Block
    def initialize
      super(0, nil, nil, nil, nil, nil)
      check_invariant
    end
  
    def add(x)
      check_invariant if $DEBUG
      n = self.next_elem
      n.prev_elem = x if n
      x.block = self
      x.next_elem = n
      x.prev_elem = self # as a sentinel
      self.next_elem = x 
      self.size += 1
      check_invariant if $DEBUG
    end
  
    def remove(x)
      check_invariant if $DEBUG
      n = x.next_elem
      x.prev_elem.next_elem = n
      n.prev_elem = x.prev_elem if n
      self.size -= 1
      check_invariant if $DEBUG
    end
  
    def each
      x = self
      while x = x.next_elem
        yield x
      end
    end
    include Enumerable

    def check_invariant
      each {|x|
        unless x.prev_elem.next_elem.equal?(x)
          raise RuntimeError
        end
      }
    end
  end
  
  UBlock = Struct.new(:size, :next_block)
  class UBlock
    def initialize
      super(0, nil)
      check_invariant
    end
  
    def add(b)
      check_invariant if $DEBUG
      b.ublock = self
      n = self.next_block
      n.prev_block = b if n
      b.next_block = n
      b.prev_block = self # as a sentinel
      self.next_block = b
      self.size += 1
      check_invariant if $DEBUG
    end
  
    def remove(b)
      check_invariant if $DEBUG
      n = b.next_block
      b.prev_block.next_block = n
      n.prev_block = b.prev_block if n
      self.size -= 1
      check_invariant if $DEBUG
    end
  
    def each
      b = self
      while b = b.next_block
        yield b
      end
    end  
    include Enumerable
  
    def complex?; size > 1; end

    def check_invariant
      each {|x|
        unless x.prev_block.next_block.equal?(x)
          raise RuntimeError
        end
      }      
    end
  end
  
  TmpList = Struct.new(:first)
  class TmpList
    SENTINEL = :sentinel
    def initialize
      super(SENTINEL)
    end
  
    def add(x)
      x.tmp_next = self.first
      self.first = x
    end
  
    def each
      x = first
      while not SENTINEL.equal?(x)
        yield x
        x = x.tmp_next
      end
    end
    include Enumerable
  
    def clear
      x = first
      while not SENTINEL.equal?(x)
        old = x
        x = x.tmp_next
        old.tmp_next = nil
      end
      self.first = SENTINEL
    end
  
    def include?(x)
      not x.tmp_next.nil?
    end
  end
  
  def initialize(g, partition)
    @ublocks = []
    @complex_ublocks = []
    @childless_blocks = []

    if partition.nil?
      b0 = []
      g.tsort_each_node {|x| b0.push x }
      partition = [b0]
    end

    h = Hash.new
    u = UBlock.new
  
    partition.each {|b0|
      b1 = Block.new
      b2 = Block.new
      u.add(b1)
      
      b0.each {|x0|
        childless = true
        g.tsort_each_child(x0) { childless = false; break }
  
        h[x0] = x = Elem.new(x0)
        if childless
          b2.add(x)
        else
          b1.add(x)
        end
      }
  
      @childless_blocks.push(b2) unless b2.size==0
    }
  
    h.each_value {|x|
      c = Count.new(0)
      g.tsort_each_child(x.data) {|y0|
        if h.has_key? y0
          y = h[y0]
          y.edges.push(Edge.new(x, c))
          c.value += 1
        end
      }
    }

    @ublocks.push(u)
    @complex_ublocks.push(u) if u.complex?

    # Preprocess childless blocks
    set = TmpList.new
    @childless_blocks.each {|b|
      b.each {|y|
        y.edges.each {|e|
          x = e.src
          set.add(x) unless set.include? x
        }
      }

      split(set)
      set.clear
    }
  end
  
  def split(e)
    splits = []

    e.each {|x|
      d = x.block
      if (assoc = d.assoc).nil?
        assoc = d.assoc = Block.new
        d.ublock.add(assoc)
        splits.push(d)
      end
      d.remove(x)
      assoc.add(x)
    }

    splits.each {|d|
      s = d.ublock
      if d.size==0
        s.remove(d)
      elsif s.size == 2
        # uのブロック数が1から2になった場合、新たにcomplexになるので
        @complex_ublocks.push(s)
      end
      d.assoc = nil
    }

    splits.clear
  end
  
  def refine(s)
    set = TmpList.new

    # step 1
    b1 = s.next_block
    b2 = b1.next_block
    b = (b1.size < b2.size) ? b1 : b2

    # step 2
    s2 = UBlock.new
    s.remove(b)
    s2.add(b)
    @ublocks.push(s2)
    @complex_ublocks.push(s) if s.complex?

    # step 3
    b_dup = []
    b.each {|y|
      b_dup.push(y)
      y.edges.each {|e|
        x = e.src
        if set.include? x
          x.count.value += 1
        else
          x.count = Count.new(1)
          set.add(x)
        end
      }
    }

    # step 4
    split(set)
    set.clear

    # step 5
    b_dup.each {|y|
      y.edges.each {|e|
        x = e.src
        if x.count.value == e.count.value and not set.include?(x)
          set.add(x)
        end
      }
    }

    # step 6
    split(set)
    set.clear

    # step 7
    b_dup.each {|y|
      y.edges.each {|e|
        x = e.src
        e.count = x.count if (e.count.value -= 1) == 0
      }
    }
    b_dup.clear

    nil
  end

  def stabilize  
    while s = @complex_ublocks.pop
      refine(s)
    end
  end

  def each_block(&block)
    @ublocks.each {|s|
      s.each(&block)
    }
    @childless_blocks.each(&block)
  end

  def self.stabilize(g, partition = nil)
    stabilizer = self.new(g, partition)
    stabilizer.stabilize

    ret = []
    stabilizer.each_block {|b|
      ret.push(b.map {|x| x.data })
    }
    ret
  end

end # class PaigeTarjan


if $0 == __FILE__
  class Hash
    alias tsort_each_node each_key
    def tsort_each_child(node, &block)
      fetch(node).each(&block)
    end
  end

  g = { 0 => [], 1 => [1], 2 => [0, 1], 3 => [1] }

  p PaigeTarjan.stabilize(g)
  #=> [[3, 1], [2], [0]]

  p PaigeTarjan.stabilize(g, [[0,1,2], [3]])
  #=> [[2], [3], [1], [0]]
end