Implement a weighted shuffle in Ruby

In my previous post I’ve implemented a Fisher-Yates shuffle algoritm in Ruby. Today I want to try somenthing complex: a weighted (or priority) based shuffle.

Given an array of hashes like this:

weighted_array = [
  {value: "HIGH 1", weight: 4},  
  {value: "HIGH 2", weight: 4},  
  {value: "MID 1", weight: 2},  
  {value: "MID 2", weight: 2},  
  {value: "LOW 1", weight: 1},  
  {value: "LOW 2", weight: 1},  
]

I want to get an array of 10 elements randomly picked across it based on their weights; something like this:

result = ["HIGH 1", "MID 2", "HIGH 2", "HIGH 1", "LOW 2", "MID 1", "HIGH 2", "HIGH 1", "MID 1", "HIGH 2"]

To do this I need to calculate the probability of any weight present in the first array. In order to achieve this I have to calculate the sums of all distinct weights.

def sum_of_weights weighted
  # .to_f for decimals in probability calculation
  weighted.map{ |w| w[:weight] }.uniq.inject(:+).to_f
end

Now I can call sum = sum_of_weights(weighted_array) and get 7, with the sum I can get the probability of each weight, simply with probability = weighted_array[0][:weight] / sum and so on

Now I can define the cut offs associated to each weight using the probability calculator below.

def probabilities weighted_array
  sum_of_weights = sum_of_weights(weighted_array)

  # initial probability
  probability = 0.0

  weights = weighted_array.map{ |w| w[:weight] }.uniq.sort
  weights.map do |w|
    current_probability = (w / sum_of_weights).round(3)
    # sum the current_probability with the previous
    probability += current_probability
    { weight: w, probability: probability }
  end
end

Given the weighted_array the above function returns an array like this:

probabilities = [
  { weight: 1, probability: 0.143 },
  { weight: 2, probability: 0.429 },
  { weight: 4, probability: 1.0 },
]

The next step is to generate a random number to select a weight based on its probability. Conceptually I need a series of if..elsif, but given the probabilities array I can use this function:

def get_weight probabilities
  r = rand
  probabilities.each do |p|
    return p[:weight] if r < p[:probability]
  end
end

Once I get the weight randomly, I can pick an element from weighted_array, filtered by the selected weight. With the ruby Array#select method:

def by_weight a, weight
  a.select{ |el| el[:weight] == weight }
end

At this point I can use the Array#sample method and get a random element from the filtered array. By repeating this 10 times I get the desidered result!

result = []
10.times do
  weight = get_weight probabilities
  result << by_weight(weighted_array, weight).sample[:value]
end

You can find all the code and little more here

During my research I have discovered various methods to achieve the desidered result. What I presented here is what I think has the best trade off between difficulty and result.

I hope to explore in more detail this sort of thing in my next post.

Bye and see you soon!

Leave a Reply

Please Login to comment

This site uses Akismet to reduce spam. Learn how your comment data is processed.