Weighted Random Array

I had an array that I was trying to do weighted random pulls from. There are a few basic algorithms for this online, but my issue was that they were not efficient for very, very large tables — most were O(n) for each time you wanted to pull an element. So I wrote my own wrapper class that would handle the task for me using Alias Tables.

class WeightedRandomArray < Array
      def initialize(other_array, weights)
        total_weights = weights.inject(0.0) { |t,e| t+e }
        proportions = weights.map { |e| e / total_weights }
 
        elements = other_array.zip(proportions)
 
        # construct an alias table for faster access
        @table = []
 
        n = elements.size
        elements.map! { |a,w| [a, w*(n-1)] }
        elements.sort! { |e1,e2| e1[1] <=> e2[1] }
 
        while elements.size > 2
          p = elements[0][1]
          elements[-1][1] -= (1.0 - p)
          @table << [p, elements[0][0], elements[-1][0]]
          elements.delete_at(0)
          elements.sort! { |a,b| a[1] <=> b[1] }
        end
 
        p = elements[0][1]
        elements[-1][1] -= p
        @table << [p, elements[0][0], elements[-1][0]]
      end
 
      def random_element
        entry = @table[(Kernel.rand * @table.size).floor]
        if Kernel.rand < entry[0]
          return entry[1]
        else
          return entry[2]
        end
      end
    end

So now the work is all up-front — and while it is considerable, so long as you will be doing enough random draws, the O(1) time to get a random entry should trump…

  • Share/Bookmark

Tags:

This entry was posted on Wednesday, June 16th, 2010 at 12:53 pmand is filed under programming. You can follow any responses to this entry through the RSS 2.0 feed. You can leave a response, or trackback from your own site.

Leave a Reply