diff --git a/README.md b/README.md index 42141a6..c457cc1 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,17 @@ raised. lume.randomchoice({true, false}) -- Returns either true or false ``` +### lume.weightedchoice(t) +Takes the argument table `t` where the keys are the possible choices and the +values are the choice's weight. A weight should be 0 or above, the larger the +weight value the higher the probability of that choice being picked. If the +table is empty or all the weights are 0 or a weight is below zero an error is +raised. +```lua +lume.weightedchoice({ ["cat"] = 10, ["dog"] = 5, ["frog"] = 0 }) +-- Returns either "cat" or "dog" with "cat" being twice as likely to be chosen. +``` + ### lume.shuffle(t) Shuffles the values of array `t` in place, returns the array. diff --git a/lume.lua b/lume.lua index 2de5867..b06980d 100644 --- a/lume.lua +++ b/lume.lua @@ -67,6 +67,21 @@ function lume.randomchoice(t) end +function lume.weightedchoice(t) + local sum = 0 + for k, v in pairs(t) do + assert(v >= 0, "weight value less than zero") + sum = sum + v + end + assert(sum ~= 0, "all weights are zero") + local rnd = lume.random(sum) + for k, v in pairs(t) do + if rnd < v then return k end + rnd = rnd - v + end +end + + function lume.shuffle(t) for i = 1, #t do local r = math.random(#t)