probabilities.rb in games_dice-0.3.0

- old
+ new

@@ -216,39 +216,118 @@
       end
     end
     GamesDice::Probabilities.new( new_probs, combined_min )
   end
 
-
   # Adds a distribution to itself repeatedly, to simulate a number of dice
   # results being summed.
-  # @param [GamesDice::Probabilities] pd Distribution to repeat
   # @param [Integer] n Number of repetitions, must be at least 1
-  # @return [GamesDice::Probabilities]
-  def self.repeat_distribution pd, n
+  # @return [GamesDice::Probabilities] new distribution
+  def repeat_sum n
     n = Integer( n )
     raise "Cannot combine probabilities less than once" if n < 1
     revbin = n.to_s(2).reverse.each_char.to_a.map { |c| c == '1' }
-    pd_power = pd
+    pd_power = self
     pd_result = nil
     max_power = revbin.count - 1
 
     revbin.each_with_index do |use_power, i|
       if use_power
         if pd_result
-          pd_result = add_distributions( pd_result, pd_power )
+          pd_result = GamesDice::Probabilities.add_distributions( pd_result, pd_power )
         else
           pd_result = pd_power
         end
       end
-      pd_power = add_distributions( pd_power, pd_power ) unless i == max_power
+      pd_power = GamesDice::Probabilities.add_distributions( pd_power, pd_power ) unless i == max_power
     end
     pd_result
   end
 
+  # Calculates distribution generated by summing best k results of n iterations
+  # of the distribution.
+  # @param [Integer] n Number of repetitions, must be at least 1
+  # @param [Integer] k Number of best results to keep and sum
+  # @return [GamesDice::Probabilities] new distribution
+  def repeat_n_sum_k n, k, kmode = :keep_best
+    n = Integer( n )
+    k = Integer( k )
+    raise "Cannot combine probabilities less than once" if n < 1
+    if k >= n
+      return repeat_sum( n )
+    end
+    new_probs = Array.new( @probs.count * k, 0.0 )
+    new_offset = @offset * k
+    d = n - k
+
+    each do | q, p_maybe |
+      next unless p_maybe > 0.0
+
+      # keep_distributions is array of Probabilities, indexed by number of keepers > q, which is in 0...k
+      keep_distributions = calc_keep_distributions( k, q, kmode )
+      p_table = calc_p_table( q, p_maybe, kmode )
+
+      (0...k).each do |n|
+        keepers = [2] * n + [1] * (k-n)
+        p_so_far = keepers.inject(1.0) { |p,idx| p * p_table[idx] }
+        next unless p_so_far > 0.0
+        (0..d).each do |dn|
+          discards = [1] * (d-dn) + [0] * dn
+          sequence = keepers + discards
+          p_sequence = discards.inject( p_so_far ) { |p,idx| p * p_table[idx] }
+          next unless p_sequence > 0.0
+          p_sequence *= GamesDice::Combinations.count_variations( sequence )
+          kd = keep_distributions[n]
+          kd.each { |r,p_r| new_probs[r-new_offset] += p_r * p_sequence }
+        end
+      end
+    end
+    GamesDice::Probabilities.new( new_probs, new_offset )
+  end
+
   private
 
+  def calc_keep_distributions k, q, kmode
+    if kmode == :keep_best
+      keep_distributions = [ GamesDice::Probabilities.new( [1.0], q * k ) ]
+      if p_gt(q) > 0.0 && k > 1
+        kd_probabilities = given_ge( q + 1 )
+        (1...k).each do |n|
+          extra_o = GamesDice::Probabilities.new( [1.0], q * ( k - n ) )
+          n_probs = kd_probabilities.repeat_sum( n )
+          keep_distributions[n] = GamesDice::Probabilities.add_distributions( extra_o, n_probs )
+        end
+      end
+    elsif kmode == :keep_worst
+      keep_distributions = [ GamesDice::Probabilities.new( [1.0], q * k ) ]
+      if p_lt(q) > 0.0 && k > 1
+        kd_probabilities = given_le( q - 1 )
+        (1...k).each do |n|
+          extra_o = GamesDice::Probabilities.new( [1.0], q * ( k - n ) )
+          n_probs = kd_probabilities.repeat_sum( n )
+          keep_distributions[n] = GamesDice::Probabilities.add_distributions( extra_o, n_probs )
+        end
+      end
+    else
+      raise "Keep mode #{kmode.inspect} not recognised"
+    end
+    keep_distributions
+  end
+
+  def calc_p_table q, p_maybe, kmode
+    if kmode == :keep_best
+      p_kept = p_gt(q)
+      p_rejected = p_lt(q)
+    elsif kmode == :keep_worst
+      p_kept = p_lt(q)
+      p_rejected = p_gt(q)
+    else
+      raise "Keep mode #{kmode.inspect} not recognised"
+    end
+    [ p_rejected, p_maybe, p_kept ]
+  end
+
   # Convert hash to array,offset notation
   def self.prob_h_to_ao h
     rmin,rmax = h.keys.minmax
     o = rmin
     a = Array.new( 1 + rmax - rmin, 0.0 )
@@ -268,5 +347,57 @@
     @probs.each_with_index { |v,i| total += (i+@offset)*v }
     total
   end
 
 end # class GamesDice::Probabilities
+
+# Helper module with optimised Ruby for counting variations of arrays, such as those returned by
+# Array#repeated_combination
+#
+# @example How many ways can [3,3,6] be arranged?
+#  GamesDice::Combinations.count_variations( [3,3,6] )
+#  => 3
+#
+# @example When prob( a ) and result( a ) are same for any arrangement of Array a
+#  items = [1,2,3,4,5,6]
+#  items.repeated_combination(5).each do |a|
+#    this_result = result( a )
+#    this_prob = prob( a ) * GamesDice::Combinations.count_variations( a )
+#    # Do something useful with this knowledge! E.g. save it to probability array.
+#  end
+#
+module GamesDice::Combinations
+  @@variations_cache = {}
+  @@factorial_cache = [1, 1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800]
+
+  # Counts variations of an array. A unique variation is an arrangement of the array elements which
+  # is detectably different (using ==) from any other. So [1,1,1] has only 1 unique arrangement,
+  # but [1,2,3] has 6 possibilities.
+  # @param [Array] array List of things that can be arranged
+  # @return [Integer] Number of unique arrangements
+  def self.count_variations array
+    all_count = array.count
+    group_sizes = group_counts( array )
+    cache_key = all_count.to_s + ":" + group_sizes.join(',')
+    @@variations_cache[cache_key] ||= variations_of( all_count, group_sizes )
+  end
+
+  private
+
+  def self.variations_of all_count, groups
+    all_arrangements = factorial( all_count )
+    # The reject is an optimisation to avoid calculating and multplying by factorial(1) (==1)
+    identical_arrangements = groups.reject {|x| x==1 }.inject(1) { |prod,g| prod * factorial(g) }
+    all_arrangements/identical_arrangements
+  end
+
+  # Returns counts of unique items in array e.g. [8,8,8,7,6,6] returns [1,2,3]
+  # Sort is for caching
+  def self.group_counts array
+    array.group_by {|x| x}.values.map {|v| v.count}.sort
+  end
+
+  def self.factorial n
+    # Can start range from 2 because we have pre-cached the result for n=1
+    @@factorial_cache[n] ||= (2..n).inject(:*)
+  end
+end