test/math/combination_test.cr

• View this file on GitHub
• Last update: 2022-03-28 10:35:50+09:00
• Problem: https://yukicoder.me/problems/no/117

Depends on

• src/math/combination.cr
• src/math/mint.cr

Code

# verification-helper: PROBLEM https://yukicoder.me/problems/no/117
require "../../src/math/mint"
require "../../src/math/combination"
C = Combination(Mint).new

read_line.to_i.times do
  read_line =~ /(.).(\d+).(\d+)./
  c, n, r = $1[-1], $2.to_i, $3.to_i
  puts case c
  when 'C'
    C.combination(n, r)
  when 'P'
    C.permutation(n, r)
  when 'H'
    C.repeated_combination(n, r)
  end
end

# verification-helper: PROBLEM https://yukicoder.me/problems/no/117
# require "../../src/math/mint"
struct ModInt(MOD)
  def self.mod
    MOD
  end

  def self.zero
    new
  end

  def self.raw(value : Int64)
    result = new
    result.value = value
    result
  end

  macro [](*nums)
    Array({{@type}}).build({{nums.size}}) do |buffer|
      {% for num, i in nums %}
        buffer[{{i}}] = {{@type}}.new({{num}})
      {% end %}
      {{nums.size}}
    end
  end

  getter value : Int64

  private macro check_mod
    {% if MOD.is_a?(NumberLiteral) %}
      {% raise "can't instantiate ModInt(MOD) with MOD = #{MOD} (MOD must be positive)" unless MOD >= 1 %}
      {% raise "can't instantiate ModInt(MOD) with MOD = #{MOD.kind} (MOD must be Int64)" unless MOD.kind == :i64 %}
    {% else %}
      {% raise "can't instantiate ModInt(MOD) with MOD = #{MOD.class_name.id} (MOD must be an integer)" %}
    {% end %}
  end

  def initialize
    check_mod
    @value = 0i64
  end

  def initialize(value)
    check_mod
    @value = value.to_i64 % MOD
  end

  def initialize(m : ModInt(MOD2)) forall MOD2
    {% raise "Can't create ModInt(#{MOD}) from ModInt(#{MOD2})" if MOD != MOD2 %}
    check_mod
    @value = m.value
  end

  protected def value=(value : Int64)
    @value = value
  end

  def self.scan(scanner, io : IO) : self
    new scanner.i64(io)
  end

  def ==(m : ModInt(MOD2)) forall MOD2
    {% raise "Can't compare ModInt(#{MOD}) and ModInt(#{MOD2})" if MOD != MOD2 %}
    value == m.value
  end

  def ==(m : Int)
    value == m
  end

  def + : self
    self
  end

  def - : self
    self.class.raw(value != 0 ? MOD &- value : 0i64)
  end

  def +(v)
    self + self.class.new(v)
  end

  def +(m : self)
    x = value &+ m.value
    x &-= MOD if x >= MOD
    self.class.raw(x)
  end

  def -(v)
    self - self.class.new(v)
  end

  def -(m : self)
    x = value &- m.value
    x &+= MOD if x < 0
    self.class.raw(x)
  end

  def *(v)
    self * self.class.new(v)
  end

  def *(m : self)
    self.class.new(value &* m.value)
  end

  def /(v)
    self / self.class.new(v)
  end

  def /(m : self)
    raise DivisionByZeroError.new if m.value == 0
    a, b, u, v = m.value, MOD, 1i64, 0i64
    while b != 0
      t = a // b
      a &-= t &* b
      a, b = b, a
      u &-= t &* v
      u, v = v, u
    end
    self.class.new(value &* u)
  end

  def //(v)
    self / v
  end

  def **(exponent : Int)
    t, res = self, self.class.raw(1i64)
    while exponent > 0
      res *= t if exponent & 1 == 1
      t *= t
      exponent >>= 1
    end
    res
  end

  {% for op in %w[< <= > >=] %}
    def {{op.id}}(other)
      raise NotImplementedError.new({{op}})
    end
  {% end %}

  def inv
    self.class.raw(1) // self
  end

  def succ
    self.class.raw(value != MOD &- 1 ? value &+ 1 : 0i64)
  end

  def pred
    self.class.raw(value != 0 ? value &- 1 : MOD &- 1)
  end

  def abs
    self
  end

  def abs2
    self * self
  end

  def to_i64 : Int64
    value
  end

  def to_s(io : IO) : Nil
    value.to_s(io)
  end

  def inspect(io : IO) : Nil
    value.inspect(io)
  end
end

struct Int
  def to_m(type : M.class) forall M
    M.new(self)
  end
end

class String
  def to_m(type : M.class) forall M
    M.new(self)
  end
end

alias Mint = ModInt(1000000007i64)
alias Mint2 = ModInt(998244353i64)

# require "../../src/math/combination"
class Combination(T)
  def initialize(initial_capacity : Int = 2)
    initial_capacity += 1
    @size = 2
    @factorial = Array(T).new(initial_capacity)
    @factorial << T.new(1) << T.new(1)
    @inv = Array(T).new(initial_capacity)
    @inv << T.zero << T.new(1)
    @finv = Array(T).new(initial_capacity)
    @finv << T.new(1) << T.new(1)
    expand_until(initial_capacity)
  end

  private def expand_until(n : Int)
    while @size <= n
      @factorial << @factorial[-1] * @size
      @inv << -@inv[T.mod % @size] * (T.mod // @size)
      @finv << @finv[-1] * @inv[@size]
      @size += 1
    end
  end

  def factorial(n : Int)
    raise IndexError.new if n < 0
    expand_until(n)
    @factorial.unsafe_fetch(n)
  end

  def inv(n : Int)
    raise DivisionByZeroError.new if n == 0
    raise IndexError.new if n < 0
    expand_until(n)
    @inv.unsafe_fetch(n)
  end

  def finv(n : Int)
    raise IndexError.new if n < 0
    expand_until(n)
    @finv.unsafe_fetch(n)
  end

  def permutation(n : Int, r : Int)
    (n < r || n < 0 || r < 0) ? T.zero : factorial(n) * finv(n - r)
  end

  def combination(n : Int, r : Int)
    (n < r || n < 0 || r < 0) ? T.zero : factorial(n) * finv(r) * finv(n - r)
  end

  def repeated_combination(n : Int, r : Int)
    (n < 0 || r < 0) ? T.zero : r == 0 ? T.new(1) : combination(n + r - 1, r)
  end

  def self.table(n : Int)
    table = Array.new(n + 1) { Array.new(n + 1, T.zero) }
    (0..n).each do |i|
      table[i][0] = table[i][i] = 1
    end
    (1..n).each do |i|
      (1...i).each do |j|
        table[i][j] = table[i - 1][j - 1] + table[i - 1][j]
      end
    end
    table
  end
end

C = Combination(Mint).new

read_line.to_i.times do
  read_line =~ /(.).(\d+).(\d+)./
  c, n, r = $1[-1], $2.to_i, $3.to_i
  puts case c
  when 'C'
    C.combination(n, r)
  when 'P'
    C.permutation(n, r)
  when 'H'
    C.repeated_combination(n, r)
  end
end

Back to top page