Kaprekar's Constant in Ruby

example.rb

require_relative 'kaprekar'

include Kaprekar

p kaprekar(four_digit_number: 9000).first(10)
#>> [9000, 8991, 8082, 8532, 6174, 6174, 6174, 6174, 6174, 6174]

p kaprekar(four_digit_number: 9000).take_while { _1 != Kaprekar::CONSTANT }
#>> [9000, 8991, 8082, 8532]

kaprekar.rb

module Kaprekar
  CONSTANT = 6174

  module Undigits
    refine Array do
      def undigits(base = 10)
        each_with_index.sum do |digit, exponent|
          digit * base**exponent
        end
      end
      
      def safe_undigits(_base = 10)
        reverse.drop_while(&:zero?).reverse.undigits
      end
    end
  end

  using Undigits

  module_function

  def kaprekar(four_digit_number:)
    raise unless four_digit_number.digits.size == 4
    
    Enumerator.produce(four_digit_number) do |number|
      sorted_digits = number.digits.sort
      
      sorted_digits.safe_undigits - sorted_digits.reverse.safe_undigits
    end
  end
end