CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
#
2
# This Module was developed by Thomas Hafner.
3
# No other references about the author.
4
#
5

6
# TITLE:
7
#
8
#   Cartesian
9
#
10
# SUMMARY:
11
#
12
#   Cartesian product and similar methods.
13
#
14
# AUTHORS:
15
#
16
#   - Thomas Hafner
17

18
#
19
module Enumerable
20

21
  class << self
22
    # Provides the cross-product of two or more Enumerables.
23
    # This is the class-level method. The instance method
24
    # calls on this.
25
    #
26
    #   Enumerable.cart([1,2], [4], ["apple", "banana"])
27
    #   #=> [[1, 4, "apple"], [1, 4, "banana"], [2, 4, "apple"], [2, 4, "banana"]]
28
    #
29
    #   Enumerable.cart([1,2], [3,4])
30
    #   #=> [[1, 3], [1, 4], [2, 3], [2, 4]]
31

32
    def cartesian_product(*enums, &block)
33
      result = [[]]
34
      while [] != enums
35
        t, result = result, []
36
        b, *enums = enums
37
        t.each do |a|
38
          b.each do |n|
39
            result << a + [n]
40
          end
41
        end
42
      end
43
      if block_given?
44
        result.each{ |e| block.call(e) }
45
      else
46
        result
47
      end
48
    end
49

50
    alias_method :cart, :cartesian_product
51
  end
52

53
  # The instance level version of <tt>Enumerable::cartesian_product</tt>.	
54
  #
55
  #   a = []
56
  #   [1,2].cart([4,5]){|elem| a << elem }
57
  #   a  #=> [[1, 4],[1, 5],[2, 4],[2, 5]]
58

59
  def cartesian_product(*enums, &block)
60
    Enumerable.cartesian_product(self, *enums, &block)
61
  end
62

63
  alias :cart :cartesian_product
64

65
  # Operator alias for cross-product.
66
  #
67
  #   a = [1,2] ** [4,5]
68
  #   a  #=> [[1, 4],[1, 5],[2, 4],[2, 5]]
69
  #
70
  def **(enum)
71
    Enumerable.cartesian_product(self, enum)
72
  end
73

74
  # Expected to be an enumeration of arrays. This method
75
  # iterates through combinations of each in position.
76
  #
77
  #   a = [ [0,1], [2,3] ]
78
  #   a.each_combo { |c| p c }
79
  #
80
  # produces
81
  #
82
  #   [0, 2]
83
  #   [0, 3]
84
  #   [1, 2]
85
  #   [1, 3]
86
  #
87
  def each_combo
88
    a = collect{ |x|
89
      x.respond_to?(:to_a) ? x.to_a : 0..x
90
    }
91

92
    if a.size == 1
93
      r = a.shift
94
      r.each{ |n|
95
        yield n
96
      }
97
    else
98
      r = a.shift
99
      r.each{ |n|
100
        a.each_combo{ |s|
101
          yield [n, *s]
102
        }
103
      }
104
    end
105
  end
106

107
  # As with each_combo but returns combos collected in an array.
108
  #
109
  def combos
110
    a = []
111
    each_combo{ |c| a << c }
112
    a
113
  end
114

115
end
116

117