weighted random distribution


See On Github

Data

Contributor

Generic placeholder thumbnail

by Yonaba

in lua

Source Code

-- Weighted/Biased random distribution implementation
-- See : http://codetheory.in/weighted-biased-random-number-generation-with-javascript-based-on-probability/

-- Note: The sum of weights should be 1. That is, all
--  weights are decimal values lower than 1.

-- Note: As all the implementations given below uses
--  math.random, one can call math.randomseed with a
--  custom seed before using them.

-- items   : an array-list of values
-- weights : a map table holding the weight for each value in
--  in the `items` list.
-- returns : an iterator function
local function expanding_random(items, weights)
  local list = {}
  for _, item in ipairs(items) do
    local n = weights[item] * 100
    for i = 1, n do table.insert(list, item) end
  end
  return function()
    return list[math.random(1, #list)]
  end
end

-- items   : an array-list of values
-- weights : a map table holding the weight for each value in
--  in the `items` list.
-- returns : an iterator function
local function in_place_random(items, weights)
  local partial_sums, partial_sum = {}, 0
  for _, item in ipairs(items) do
    partial_sum = partial_sum + weights[item]
    table.insert(partial_sums, partial_sum)
  end
  return function()
    local n, s = math.random(), 0
    for i, si in ipairs(partial_sums) do
      s = s + si
      if si > n then return items[i] end
    end
  end
end

return {
  expanding = expanding_random,
  in_place  = in_place_random
}
-- Tests for weighted_random.lua
local rand = require 'weighted_random'

-- Generates n random values
local function gen(f, items, weights, n)
  local rnd = f(items, weights)
  local t = {}
  for i = 1, n do
    local val = rnd()
    t[val] = (t[val] or 0) + 1
  end
  for k,v in pairs(t) do
    print((' >> %s : %d (%.2f %%)'):format(k, v, 100*v/n))
  end
end

print(('-'):rep(80))
local items = {'a', 'b', 'c', 'd', 'e'}
local weights = {a = 0.1, b = 0.2, c = 0.3, d = 0.2, e = 0.2}
print('In-place random distribution')
gen(rand.in_place, items, weights, 1e6)
print('\nExpanding random distribution')
gen(rand.expanding, items, weights, 1e6)
print(('-'):rep(80))