test/int/each_combination_test.cr

• View this file on GitHub
• Last update: 2021-10-17 21:56:08+09:00
• Problem: https://onlinejudge.u-aizu.ac.jp/courses/lesson/8/ITP2/all/ITP2_11_D

Depends on

• src/int/each_combination.cr

Code

# verification-helper: PROBLEM https://onlinejudge.u-aizu.ac.jp/courses/lesson/8/ITP2/all/ITP2_11_D
require "../../src/int/each_combination"
n, k = read_line.split.map(&.to_i)

ans = [] of Int32
Int.each_combination(n, k) { |x| ans << x }
ans2 = Int.each_combination(n, k).to_a
raise "" unless ans == ans2

ans.each do |x|
  print x, ": ", (0...n).select { |i| x.bit(i) == 1 }.join(' '), '\n'
end

# verification-helper: PROBLEM https://onlinejudge.u-aizu.ac.jp/courses/lesson/8/ITP2/all/ITP2_11_D
# require "../../src/int/each_combination"
struct Int
  private class CombinationIterator(T, U)
    include Iterator(T)

    @combination : T

    def initialize(@n : T, @k : U)
      @combination = (T.new(1) << @k) - 1
      @first = true
    end

    def next
      if @first
        @first = false
        @combination
      else
        x = @combination & -@combination
        y = @combination + x
        @combination = ((@combination & ~y) // x >> 1) | y
        @combination < (T.new(1) << @n) ? @combination : stop
      end
    end
  end

  # Returns an iterator that returns all integers whose bit_length is *n* and popcount is *k*.
  #
  # ```
  # Int.each_combination(3, 2).to_a # => [0b011, 0b101, 0b110]
  # ```
  def self.each_combination(n : Int, k : Int)
    CombinationIterator.new(n, k)
  end

  # Calls the given block for each integer whose bit_length is *n* and popcount is *k*.
  #
  # ```
  # # x = 0b011, 0b101, 0b110
  # Int.each_combination(3, 2) { |x| }
  # ```
  def self.each_combination(n : Int, k : Int, &) : Nil
    combination = (n.class.new(1) << k) - 1
    while combination < (n.class.new(1) << n)
      yield combination
      x = combination & (-combination)
      y = combination + x
      combination = ((combination & ~y) // x >> 1) | y
    end
  end
end

n, k = read_line.split.map(&.to_i)

ans = [] of Int32
Int.each_combination(n, k) { |x| ans << x }
ans2 = Int.each_combination(n, k).to_a
raise "" unless ans == ans2

ans.each do |x|
  print x, ": ", (0...n).select { |i| x.bit(i) == 1 }.join(' '), '\n'
end

Back to top page