From 0e34a7431fb5eb069e61bec7b1a7bc92ac17f046 Mon Sep 17 00:00:00 2001 From: Ahmed Alparslan Ozdemir Date: Thu, 16 Feb 2023 17:47:51 +0300 Subject: [PATCH] performance optimization --- weightedRandom.go | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/weightedRandom.go b/weightedRandom.go index 835a7aa..21cd193 100644 --- a/weightedRandom.go +++ b/weightedRandom.go @@ -14,14 +14,24 @@ func WeightedRandom[P interface{ GetWeight() int64 }](elements []P) (P, error) { total.Add(total, big.NewInt(element.GetWeight())) } - var shuffler []P + var rng, _ = rand.Int(rand.Reader, total) + var randomNumber = rng.Int64() + var countWeight int64 = 0 + var target P for _, element := range elements { - for j := int64(0); j < element.GetWeight(); j++ { - shuffler = append(shuffler, element) + var weight = element.GetWeight() + + if weight < 0 { + continue + } + + countWeight += weight + if randomNumber < countWeight { + target = element + break } } - var rnd, _ = rand.Int(rand.Reader, total) - return shuffler[rnd.Int64()], nil + return target, nil }