WeightedRandom - A weighted random module with a built-in luck system

A weighted random module with a built-in luck system. It is very easy to setup and use (see code example below). It has other functions than the code example below that you can use to improve your UI/UX design.

I’m releasing this early as part of my future utilities package. It is a very popular topic nowadays, and it is a very common thread in the #help-and-feedback:scripting-support category, but I don’t seem to see any publicly available, ready to use module out there.

To those developers who wants to read a little bit more into the types

The entire module is strictly typed and uses a generic as the items. If you pass a string to AddItem, all T items from that point onward will refer to as a string. Otherwise, you can strictly define it from the start using:

--!strict

local WeightedRandom = require(path)
type WeightedRandom<T> = WeightedRandom.WeightedRandom<T>

-- myWeightedRandom will now refer to all T items as an Instance
local myWeightedRandom: WeightedRandom<Instance> = WeightedRandom.new()

-- on strict mode, this will warn you
myWeightedRandom:AddItem("item", 10)

WeightedRandom.lua
--!strict

local WeightedRandom = {}
local WeightedRandomProxy: any = newproxy()

export type WeightedRandomItem = {
	Weight: number,
	LuckInfluence: number
}

export type WeightedRandomClass<T> = {
	Random: Random,
	GetItems: (self: WeightedRandom<T>) -> {[T]: WeightedRandomItem},
	AddItem: (self: WeightedRandom<T>, item: T, weight: number, luckInfluence: number?) -> (),
	RemoveItem: (self: WeightedRandom<T>, item: T) -> (),
	GetWeight: (self: WeightedRandom<T>, item: T, luckFactor: number?) -> number?,
	GetWeights: (self: WeightedRandom<T>, luckFactor: number?) -> {[T]: number},
	GetTotalWeight: (self: WeightedRandom<T>, luckFactor: number?) -> number,
	GetProbability: (self: WeightedRandom<T>, item: T, luckFactor: number?) -> number?,
	GetProbabilities: (self: WeightedRandom<T>, luckFactor: number?) -> {[T]: number},
	Next: (self: WeightedRandom<T>, luckFactor: number?) -> T
}

export type WeightedRandomMetatable<T> = {
	_internal: {
		items: {[T]: WeightedRandomItem}
	},
	_proxy: any,
	__newindex: (self: WeightedRandom<T>, index: any, value: any) -> (),
	__tostring: (self: WeightedRandom<T>) -> string
}

export type WeightedRandom<T> = typeof(setmetatable({} :: WeightedRandomClass<T>, {} :: WeightedRandomMetatable<T>))

local function clone<T>(t: T, deep: boolean?): T
	t = table.clone(t :: any) :: any

	if deep == true then
		for index: any, value: any in t :: any do
			if typeof(value) == "table" then
				(t :: any)[index] = clone(value, true)
			end
		end
	end

	return t
end

function WeightedRandom.new<T>(random: Random?): WeightedRandom<T>
	if random ~= nil and typeof(random) ~= "Random" then
		error(`invalid argument 'random' (Random expected, got {typeof(random)})`, 2)
	end
	
	local weightedRandomClass: WeightedRandomClass<T> = {
		Random = random or Random.new(),
		GetItems = function(weightedRandom: WeightedRandom<T>): {[T]: WeightedRandomItem}
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			local weightedRandomMetatable: WeightedRandomMetatable<T> = getmetatable(weightedRandom)
			
			return clone(weightedRandomMetatable._internal.items, true)
		end,
		AddItem = function(weightedRandom: WeightedRandom<T>, item: T, weight: number, luckInfluence: number?): ()
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			local weightedRandomMetatable: WeightedRandomMetatable<T> = getmetatable(weightedRandom)
			
			if item == nil then
				error(`invalid argument 'item' (must not be nil)`, 2)
			else
				local itemToCompare: T? = next(weightedRandomMetatable._internal.items)
				
				if itemToCompare ~= nil and typeof(item) ~= typeof(itemToCompare) then
					error(`invalid argument 'item' ({typeof(itemToCompare)} expected, got {typeof(item)})`, 2)
				end
			end
			
			if typeof(weight) ~= "number" then
				error(`invalid argument 'weight' (number expected, got {typeof(weight)})`, 2)
			end
			
			if luckInfluence ~= nil and typeof(luckInfluence) ~= "number" then
				error(`invalid argument 'luckInfluence' (number expected, got {typeof(luckInfluence)})`, 2)
			end
			
			weightedRandomMetatable._internal.items[item] = {
				Weight = weight,
				LuckInfluence = if luckInfluence ~= nil then luckInfluence else 0
			}
		end,
		RemoveItem = function(weightedRandom: WeightedRandom<T>, item: T): ()
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if item == nil then
				error(`invalid argument 'item' (must not be nil)`, 2)
			end
			
			local weightedRandomMetatable: WeightedRandomMetatable<T> = getmetatable(weightedRandom)
			
			weightedRandomMetatable._internal.items[item] = nil
		end,
		GetWeight = function(weightedRandom: WeightedRandom<T>, item: T, luckFactor: number?): number?
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if item == nil then
				error(`invalid argument 'item' (must not be nil)`, 2)
			end
			
			if luckFactor ~= nil and typeof(luckFactor) ~= "number" then
				error(`invalid argument 'luckFactor' (number expected, got {typeof(luckFactor)})`, 2)
			end
			
			local weightedRandomMetatable: WeightedRandomMetatable<T> = getmetatable(weightedRandom)
			
			if weightedRandomMetatable._internal.items[item] ~= nil then
				return weightedRandomMetatable._internal.items[item].Weight + (weightedRandomMetatable._internal.items[item].Weight * (weightedRandomMetatable._internal.items[item].LuckInfluence * (if luckFactor ~= nil then luckFactor else 0)))
			else
				return nil
			end
		end,
		GetWeights = function(weightedRandom: WeightedRandom<T>, luckFactor: number?): {[T]: number}
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if luckFactor ~= nil and typeof(luckFactor) ~= "number" then
				error(`invalid argument 'luckFactor' (number expected, got {typeof(luckFactor)})`, 2)
			end
			
			local weightedRandomMetatable: WeightedRandomMetatable<T> = getmetatable(weightedRandom)
			local weights: {[T]: number} = {}

			for item: T in weightedRandomMetatable._internal.items do
				weights[item] = weightedRandom:GetWeight(item, luckFactor) :: number
			end

			return weights
		end,
		GetTotalWeight = function(weightedRandom: WeightedRandom<T>, luckFactor: number?): number
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if luckFactor ~= nil and typeof(luckFactor) ~= "number" then
				error(`invalid argument 'luckFactor' (number expected, got {typeof(luckFactor)})`, 2)
			end
			
			local totalWeight: number = 0

			for _: T, weight: number in weightedRandom:GetWeights(luckFactor) do
				totalWeight += weight
			end

			return totalWeight
		end,
		GetProbability = function(weightedRandom: WeightedRandom<T>, item: T, luckFactor: number?): number?
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if item == nil then
				error(`invalid argument 'item' (must not be nil)`, 2)
			end
			
			if luckFactor ~= nil and typeof(luckFactor) ~= "number" then
				error(`invalid argument 'luckFactor' (number expected, got {typeof(luckFactor)})`, 2)
			end
			
			local weightedRandomMetatable: WeightedRandomMetatable<T> = getmetatable(weightedRandom)
			
			if weightedRandomMetatable._internal.items[item] ~= nil then
				return (weightedRandom:GetWeight(item, luckFactor) :: number) / weightedRandom:GetTotalWeight(luckFactor)
			else
				return nil
			end
		end,
		GetProbabilities = function(weightedRandom: WeightedRandom<T>, luckFactor: number?): {[T]: number}
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if luckFactor ~= nil and typeof(luckFactor) ~= "number" then
				error(`invalid argument 'luckFactor' (number expected, got {typeof(luckFactor)})`, 2)
			end
			
			local probabilities: {[T]: number} = {}

			for item: T in weightedRandom:GetWeights(luckFactor) do
				probabilities[item] = weightedRandom:GetProbability(item, luckFactor) :: number
			end

			return probabilities
		end,
		Next = function(weightedRandom: WeightedRandom<T>, luckFactor: number?): T
			if WeightedRandom.is(weightedRandom) == false then
				error(`expected to call with ':' not '.'`, 2)
			end
			
			if luckFactor ~= nil and typeof(luckFactor) ~= "number" then
				error(`invalid argument 'luckFactor' (number expected, got {typeof(luckFactor)})`, 2)
			end
			
			local weightedItems: {[T]: number} = weightedRandom:GetWeights(luckFactor)
			local totalWeight: number = weightedRandom:GetTotalWeight(luckFactor)
			local randomWeight: number = weightedRandom.Random:NextNumber() * totalWeight
			local lastItem: T = nil

			for item: T, weight: number in weightedItems do
				randomWeight -= weight
				lastItem = item

				if randomWeight <= 0 then
					return item
				end
			end

			return lastItem
		end,
	}
	
	local weightedRandomMetatable: WeightedRandomMetatable<T> = {
		_internal = {
			items = {}
		},
		_proxy = WeightedRandomProxy,
		__newindex = function(_: WeightedRandom<T>, index: any): ()
			error(`{index} cannot be assigned to`, 2)
		end,
		__tostring = function(): string
			return `WeightedRandom`
		end
	}
	
	return table.freeze(setmetatable(weightedRandomClass, table.freeze(weightedRandomMetatable)))
end

function WeightedRandom.is(any: any): boolean
	return getmetatable(any) ~= nil and getmetatable(any)._proxy == WeightedRandomProxy
end

return table.freeze(WeightedRandom)
Usage and Examples
-- setup world one pets data and weightedRandom
local world_one_pets = WeightedRandom.new()
local world_one_pets_data: {
	[string]: {
		Weight: number,
		LuckInfluence: number
	}
} = {
	["Common"] = {
		Weight = 40,
		LuckInfluence = 0
	},
	["Uncommon"] = {
		Weight = 30,
		LuckInfluence = 0
	},
	["Rare"] = {
		Weight = 15,
		LuckInfluence = 0.5
	},
	["Epic"] = {
		Weight = 8,
		LuckInfluence = 1
	},
	["Legendary"] = {
		Weight = 4,
		LuckInfluence = 2
	},
	["Mythical"] = {
		Weight = 2,
		LuckInfluence = 4
	},
	["Secret"] = {
		Weight = 1,
		LuckInfluence = 8
	}
}

for rarity, data in world_one_pets_data do
	world_one_pets:AddItem(rarity, data.Weight, data.LuckInfluence)
end

-- output total weight with no luck factor
print(`Total weight: {world_one_pets:GetTotalWeight()}`)
print()

-- output probabilities
local function outputProbabilities(luckFactor)
	print(`Probabilities with a LuckFactor of {luckFactor}`)
	
	for rarity: string, probability: number in world_one_pets:GetProbabilities(luckFactor) do
		print(`{rarity}: {string.format("%.2f", probability * 100)}%`)
	end
	
	print()
end

-- 0x luck
outputProbabilities(0)
-- 5x luck
outputProbabilities(5)
-- 100x luck
outputProbabilities(100)

-- generate 5 items with a 5x luck factor
print(`Generating 5 pets with a LuckFactor of 5 (results may vary)`)
for i = 1, 5 do
	local chosenRarity = world_one_pets:Next(5)
	
	print(`Unlocked a {string.lower(chosenRarity)} pet ({string.format("%.2f", world_one_pets:GetProbability(chosenRarity, 5) :: number * 100)}%)`)
end

Screenshot of the output from the code example above:

10 Likes

Looks pretty alright so far from this post, I’ll try it out in Studio to test it!!

Cool, let me know if you have any questions or concerns :slight_smile:

Probably the best RNG module designed for this

One question, do you know how i can add instances to the pool and also specify these with their custom weight, luck influence?

I don’t recommend mixing different types of values together, if you’re going to do Instances, you should only keep the items as an Instance (read the developer note above).

However, if you really wanted to mix for some reason, declare it as WeightedRandom<any> and remove the condition that checks for item type equality in the AddItem function (keeping only the nil check)

3 Likes