Weighted random drop table

I wrote this small weighted random function for a project I’m working on. It turned out to be general enough I figured others could use it for their own projects. The random function itself is quoted here, and there is a DefoldCodePad link to an example program where you can adjust the drop chances in the example table. It should be simple to adopt this code to your projects.

Click here for example

  • Returns a random item from item_table parameter based on the “weight” field
  • Weights are relative to each other so they do not have to sum to a certain number
  • Excluding the weight field, setting it to 0, or nil will exclude it from being picked
  • Setting all weights to 0 returns the first value from the pairs iterator
  • If no weight fields are found in the item_table values, this function will return nil
  • Use any random function with interval of [0,1) (or [0,1] should also work with the code below)
local function rand_weighted(item_table)
	local weight_sum = 0
	for _, item in pairs(item_table) do
		if item.weight then
			weight_sum = weight_sum + item.weight
		end
	end

	local r = math.random() * weight_sum
	for _, item in pairs(item_table) do
		if item.weight then
			r = r - item.weight
			if r <= 0 then
				return item
			end
		end
	end
end

The rest of the code in the Codepad example is to demonstrate the use of this. It’s very verbose in the comments and code for clariy. This is the first time I’ve contributed code to a public audience and I have no professional development experience, so please give any constructive feedback on how to better share stuff like this or future improvements.

This did not seem worth creating a git repo and an asset portal submission since it is a simple well-known algorithm, so I’ve put it together in the Defold Codepad. I also wanted a reason to use the Defold Codepad.

The weights in the example initially give the red tiles/items the biggest chance of dropping. Those can be adjusted per the instructions in the comments (there is a table at the top of the file with a big comment block). I found the Codepad limited in being able to demonstrate what I wanted to demo once I started writing it, but it is still useful to play around with.

Edit: typos, add description from example code comments, fix comment indention in Codepad link at factory.create(), clearer Codepad link

13 Likes

If the item_table has too many entries and there is the need to sample too often then one can encounter a performance bottleneck. For these cases, after some google-fu and chat-gpt’s help, I got to the much more complex(in part because of Lua’s lack of a standard queue implementation) solution:
codepad

-- Function to perform weighted choice using the Alias Method
-- weights: a table where each index represents a choice and the corresponding value represents its weight
-- Returns: the chosen item
local function newNonUniformRandomSelector(weights)
	-- Count the number of choices
	local numChoices = #weights

	-- Normalize the weights
	local sum = 0
	for i = 1, numChoices do
		sum = sum + weights[i]
	end

	for i = 1, numChoices do
		weights[i] = weights[i] / sum
	end

	-- Create the alias table and probability table
	local alias = {}
	local probability = {}
	local average = 1 / numChoices

	local small = Queue.new(numChoices)
	local large = Queue.new(numChoices)

	-- Initialize the small and large stacks
	for i = 1, numChoices do
		if weights[i] < average then
			small:enqueue(i)
		else
			large:enqueue(i)
		end
	end

	-- Build the alias table and probability table
	while not small:isEmpty() and not large:isEmpty() do
		local smallIndex = small:dequeue()
		local largeIndex = large:dequeue()

		probability[smallIndex] = weights[smallIndex] * numChoices
		alias[smallIndex] = largeIndex
		weights[largeIndex] = (weights[largeIndex] + weights[smallIndex]) - average

		if weights[largeIndex] < average then
			small:enqueue(largeIndex)
		else
			large:enqueue(largeIndex)
		end
	end

	while not large:isEmpty() do
		local largeIndex = large:dequeue()
		probability[largeIndex] = 1
	end

	while not small:isEmpty() do
		local smallIndex = small:dequeue()
		probability[smallIndex] = 1
	end

	return function()
		-- Perform the weighted choice
		local randomIndex = math.random(numChoices)
		local randomProbability = math.random()

		if randomProbability < probability[randomIndex] then
			return randomIndex
		else
			return alias[randomIndex]
		end
	end
end

local function newRandomItemSelector(item_table)
	local item_array = {}
	local weights = {}
	for _, item in pairs(item_table) do
		item_array[#item_array + 1] = item
		weights[#weights + 1] = item.weight
	end

	local weight_sum = 0
	for _, item in pairs(item_table) do
		if item.weight then
			weight_sum = weight_sum + item.weight
		end
	end

	local selector = newNonUniformRandomSelector(weights)

	return function()
		return item_array[selector()]
	end
end

It’s based on:

It requires O(n) time to build the “sampler” function but the subsequent queries are O(1). It also requires O(n) space.

4 Likes