ac-library-rb 0.5.2 → 0.6.1

data/lib/crt.rb

@@ -1,6 +1,8 @@
 # return [rem, mod] or [0, 0] (if no solution)
 def crt(r, m)
-  raise ArgumentError if r.size != m.size
+  unless r.size == m.size
+    raise ArgumentError.new("size of r and m must be equal for crt(r, m)")
+  end
 
   n = r.size
   r0, m0 = 0, 1

data/lib/dsu.rb

@@ -1,12 +1,13 @@
 # Disjoint Set Union
 class DSU
-  def initialize(n = 0)
+  def initialize(n)
+    @n = n
+    @parent_or_size = Array.new(n, -1)
     # root node: -1 * component size
     # otherwise: parent
-    @parent_or_size = Array.new(n, -1)
   end
 
-  attr_accessor :parent_or_size
+  attr_reader :parent_or_size, :n
 
   def merge(a, b)
     x = leader(a)
@@ -19,12 +20,16 @@ class DSU
   end
   alias unite merge
 
-  def same(a, b)
+  def same?(a, b)
     leader(a) == leader(b)
   end
-  alias same? same
+  alias same same?
 
   def leader(a)
+    unless 0 <= a && a < @n
+      raise ArgumentError.new, "#{a} is out of range (0...#{@n})"
+    end
+
     @parent_or_size[a] < 0 ? a : (@parent_or_size[a] = leader(@parent_or_size[a]))
   end
   alias root leader
@@ -39,6 +44,4 @@ class DSU
   end
 end
 
-UnionFind        = DSU
-UnionFindTree    = DSU
-DisjointSetUnion = DSU
+UnionFind = DSU

data/lib/fenwick_tree.rb

@@ -2,7 +2,7 @@
 class FenwickTree
   attr_reader :data, :size
 
-  def initialize(arg = 0)
+  def initialize(arg)
     case arg
     when Array
       @size = arg.size
@@ -17,7 +17,7 @@ class FenwickTree
       @size = arg
       @data = Array.new(@size + 1, 0)
     else
-      raise ArgumentError
+      raise ArgumentError.new("wrong argument. type is Array or Integer")
     end
   end
 
@@ -29,8 +29,25 @@ class FenwickTree
     end
   end
 
-  def sum(l, r)
-    _sum(r) - _sum(l)
+  # .sum(l, r)  # [l, r)  <- Original
+  # .sum(r)     # [0, r)  <- [Experimental]
+  # .sum(l..r)  # [l, r]  <- [Experimental]
+  def sum(a, b = nil)
+    if b
+      _sum(b) - _sum(a)
+    elsif a.is_a?(Range)
+      l = a.begin
+      l += @size if l < 0
+      if r = a.end
+        r += @size if r < 0
+        r += 1 unless a.exclude_end?
+      else
+        r = @size
+      end
+      _sum(r) - _sum(l)
+    else
+      _sum(a)
+    end
   end
 
   def _sum(i)
@@ -41,8 +58,5 @@ class FenwickTree
     end
     res
   end
+  alias left_sum _sum
 end
-
-FeTree            = FenwickTree
-Fetree            = FenwickTree
-BinaryIndexedTree = FenwickTree

data/lib/floor_sum.rb

@@ -1,21 +1,44 @@
 def floor_sum(n, m, a, b)
+  raise ArgumentError if n < 0 || m < 1
+
   res = 0
 
-  if a >= m
-    res += (n - 1) * n * (a / m) / 2
-    a %= m
+  if a < 0
+    a2 = a % m
+    res -= n * (n - 1) / 2 * ((a2 - a) / m)
+    a = a2
   end
 
-  if b >= m
-    res += n * (b / m)
-    b %= m
+  if b < 0
+    b2 = b % m
+    res -= n * ((b2 - b) / m)
+    b = b2
   end
 
-  y_max = (a * n + b) / m
+  res + floor_sum_unsigned(n, m, a, b)
+end
+
+def floor_sum_unsigned(n, m, a, b)
+  res = 0
+
+  while true
+    if a >= m
+      res += n * (n - 1) / 2 * (a / m)
+      a %= m
+    end
+
+    if b >= m
+      res += n * (b / m)
+      b %= m
+    end
 
-  return res if y_max == 0
+    y_max = a * n + b
+    break if y_max < m
+
+    n = y_max / m
+    b = y_max % m
+    m, a = a, m
+  end
 
-  x_max = (m * y_max - b + a - 1) / a
-  res += (n - x_max) * y_max + floor_sum(y_max, a, m, a * x_max - m * y_max + b)
   res
 end

data/lib/lazy_segtree.rb

@@ -49,7 +49,18 @@ class LazySegtree
   end
   alias [] get
 
-  def prod(l, r)
+  def prod(l, r = nil)
+    if r.nil? # if 1st argument l is Range
+      if r = l.end
+        r += @n if r < 0
+        r += 1 unless l.exclude_end?
+      else
+        r = @n
+      end
+      l = l.begin
+      l += @n if l < 0
+    end
+
     return @e if l == r
 
     l += @size
@@ -83,9 +94,22 @@ class LazySegtree
   end
 
   # apply(pos, f)
-  # apply(l, r, f) -> range_apply(l, r, f)
+  # apply(l, r, f)  -> range_apply(l, r, f)
+  # apply(l...r, f) -> range_apply(l, r, f)  ... [Experimental]
   def apply(pos, f, fr = nil)
-    return range_apply(pos, f, fr) if fr
+    if fr
+      return range_apply(pos, f, fr)
+    elsif pos.is_a?(Range)
+      l = pos.begin
+      l += @n if l < 0
+      if r = pos.end
+        r += @n if r < 0
+        r += 1 unless pos.exclude_end?
+      else
+        r = @n
+      end
+      return range_apply(l, r, f)
+    end
 
     pos += @size
     @log.downto(1) { |i| push(pos >> i) }
@@ -194,3 +218,5 @@ class LazySegtree
     @lz[k] = @id
   end
 end
+
+LazySegTree = LazySegtree

data/lib/max_flow.rb

@@ -1,6 +1,6 @@
 # MaxFlowGraph
 class MaxFlow
-  def initialize(n = 0)
+  def initialize(n)
     @n   = n
     @pos = []
     @g   = Array.new(n) { [] }
@@ -20,6 +20,15 @@ class MaxFlow
     edge_number
   end
 
+  def add_edges(edges)
+    edges.each{ |from, to, cap| add_edge(from, to, cap) }
+    self
+  end
+
+  def add(x, to = nil, cap = nil)
+    cap ? add_edge(x, to, cap) : add_edges(x)
+  end
+
   def push(edge)
     add_edge(*edge)
   end
@@ -132,6 +141,3 @@ class MaxFlow
     res
   end
 end
-
-MaxFlowGraph = MaxFlow
-MFGraph      = MaxFlow

data/lib/min_cost_flow.rb

@@ -11,7 +11,7 @@ class MinCostFlow
     @g_cost = Array.new(n) { [] }
     @pv     = Array.new(n)
     @pe     = Array.new(n)
-    @dual   = Array.new(n) { 0 }
+    @dual   = Array.new(n, 0)
   end
 
   def add_edge(from, to, cap, cost)
@@ -36,6 +36,15 @@ class MinCostFlow
     edge_number
   end
 
+  def add_edges(edges)
+    edges.each{ |from, to, cap, cost| add_edge(from, to, cap, cost) }
+    self
+  end
+
+  def add(x, to = nil, cap = nil, cost = nil)
+    cost ? add_edge(x, to, cap, cost) : add_edges(x)
+  end
+
   def get_edge(i)
     from, id = @pos[i]
     to = @g_to[from][id]
@@ -104,7 +113,7 @@ class MinCostFlow
   def slope(s, t, flow_limit = Float::MAX)
     flow = 0
     cost = 0
-    prev_cost = -1
+    prev_cost_per_flow = -1
     result = [[flow, cost]]
 
     while flow < flow_limit
@@ -129,15 +138,12 @@ class MinCostFlow
       d = -@dual[s]
       flow += c
       cost += c * d
-      result.pop if prev_cost == d
+      result.pop if prev_cost_per_flow == d
       result << [flow, cost]
-      prev_cost = d
+      prev_cost_per_flow = d
     end
 
     result
   end
   alias min_cost_slop slope
 end
-
-MCF      = MinCostFlow
-MCFGraph = MinCostFlow

data/lib/scc.rb

@@ -1,15 +1,29 @@
 # Strongly Connected Components
-class SCCGraph
+class SCC
   # initialize graph with n vertices
-  def initialize(n = 0)
-    @n, @edges = n, []
+  def initialize(n)
+    @n = n
+    @edges = []
   end
 
   # add directed edge
   def add_edge(from, to)
-    raise "invalid params" unless (0...@n).include? from and (0...@n).include? to
+    unless 0 <= from && from < @n and 0 <= to && to < @n
+      msg = "Wrong params: from=#{from} and to=#{to} must be in 0...#{@n}"
+      raise ArgumentError.new(msg)
+    end
 
     @edges << [from, to]
+    self
+  end
+
+  def add_edges(edges)
+    edges.each{ |from, to| add_edge(from, to) }
+    self
+  end
+
+  def add(x, to = nil)
+    to ? add_edge(x, to) : add_edges(x)
   end
 
   # returns list of strongly connected components
@@ -17,10 +31,8 @@ class SCCGraph
   # O(@n + @edges.size)
   def scc
     group_num, ids = scc_ids
-    counts = [0] * group_num
-    ids.each { |x| counts[x] += 1 }
     groups = Array.new(group_num) { [] }
-    ids.each_with_index { |x, i| groups[x] << i }
+    ids.each_with_index { |id, i| groups[id] << i }
     groups
   end
 
@@ -59,7 +71,8 @@ class SCCGraph
   end
 
   def csr
-    start, elist = [0] * (@n + 1), [nil] * @edges.size
+    start = [0] * (@n + 1)
+    elist = [nil] * @edges.size
     @edges.each { |(i, _)| start[i + 1] += 1 }
     @n.times { |i| start[i + 1] += start[i] }
     counter = start.dup
@@ -70,8 +83,3 @@ class SCCGraph
     [start, elist]
   end
 end
-
-# class alias
-StronglyConnectedComponents = SCCGraph
-SCC  = SCCGraph
-SCCG = SCCGraph

data/lib/segtree.rb

@@ -2,14 +2,10 @@
 class Segtree
   attr_reader :d, :op, :n, :leaf_size, :log
 
-  # new(e){ |x, y|  }
   # new(v, e){ |x, y|  }
   # new(v, op, e)
-  def initialize(a0, a1 = nil, a2 = nil, &block)
-    if a1.nil?
-      @e, @op = a0, proc(&block)
-      v = []
-    elsif a2.nil?
+  def initialize(a0, a1, a2 = nil, &block)
+    if a2.nil?
       @e, @op = a1, proc(&block)
       v = (a0.is_a?(Array) ? a0 : [@e] * a0)
     else
@@ -37,7 +33,18 @@ class Segtree
   end
   alias [] get
 
-  def prod(l, r)
+  def prod(l, r = nil)
+    if r.nil? # if 1st argument l is Range
+      if r = l.end
+        r += @n if r < 0
+        r += 1 unless l.exclude_end?
+      else
+        r = @n
+      end
+      l = l.begin
+      l += @n if l < 0
+    end
+
     return @e if l == r
 
     sml = @e
@@ -127,20 +134,19 @@ class Segtree
     @d[k] = @op.call(@d[2 * k], @d[2 * k + 1])
   end
 
-  def inspect
-    t = 0
-    res = "SegmentTree @e = #{@e}, @n = #{@n}, @leaf_size = #{@leaf_size} @op = #{@op}\n  "
-    a = @d[1, @d.size - 1]
-    a.each_with_index do |e, i|
-      res << e.to_s << ' '
-      if t == i && i < @leaf_size
-        res << "\n  "
-        t = t * 2 + 2
-      end
-    end
-    res
-  end
+  # def inspect # for debug
+  #   t = 0
+  #   res = "Segtree @e = #{@e}, @n = #{@n}, @leaf_size = #{@leaf_size} @op = #{@op}\n  "
+  #   a = @d[1, @d.size - 1]
+  #   a.each_with_index do |e, i|
+  #     res << e.to_s << ' '
+  #     if t == i && i < @leaf_size
+  #       res << "\n  "
+  #       t = t * 2 + 2
+  #     end
+  #   end
+  #   res
+  # end
 end
 
-SegTree     = Segtree
-SegmentTree = Segtree
+SegTree = Segtree

data/lib/two_sat.rb

@@ -3,23 +3,25 @@ require_relative './scc.rb'
 # TwoSAT
 # Reference: https://github.com/atcoder/ac-library/blob/master/atcoder/twosat.hpp
 class TwoSAT
-  def initialize(n = 0)
+  def initialize(n)
     @n = n
     @answer = Array.new(n)
-    @scc = SCCGraph.new(2 * n)
+    @scc = SCC.new(2 * n)
   end
 
   attr_reader :answer
 
   def add_clause(i, f, j, g)
-    raise RangeError unless (0...@n).cover?(i) && (0...@n).cover?(j)
+    unless 0 <= i && i < @n and 0 <= j && j < @n
+      raise ArgumentError.new("i:#{i} and j:#{j} must be in (0...#{@n})")
+    end
 
     @scc.add_edge(2 * i + (f ? 0 : 1), 2 * j + (g ? 1 : 0))
     @scc.add_edge(2 * j + (g ? 0 : 1), 2 * i + (f ? 1 : 0))
     nil
   end
 
-  def satisfiable
+  def satisfiable?
     id = @scc.send(:scc_ids)[1]
     @n.times do |i|
       return false if id[2 * i] == id[2 * i + 1]
@@ -28,7 +30,7 @@ class TwoSAT
     end
     true
   end
-  alias satisfiable? satisfiable
+  alias satisfiable satisfiable?
 end
 
 TwoSat = TwoSAT