diff --git a/src/GameUtils/Extensions/CollectionExtensions.cs b/src/GameUtils/Extensions/CollectionExtensions.cs index 333b61e..50daa64 100644 --- a/src/GameUtils/Extensions/CollectionExtensions.cs +++ b/src/GameUtils/Extensions/CollectionExtensions.cs @@ -274,37 +274,34 @@ public static T WeightedRandom(this IEnumerable source, Func wei return list[count - 1]; } - var array = source.ToArray(); - int arrayCount = array.Length; - if (arrayCount == 0) - { - throw new InvalidOperationException("Sequence contains no elements."); - } + T selected = default!; + float fallbackTotalWeight = 0f; + bool hasElements = false; - var arrayTotalWeight = 0f; - for (int i = 0; i < arrayCount; i++) + foreach (var item in source) { - arrayTotalWeight += weightSelector(array[i]); + hasElements = true; + float weight = weightSelector(item); + if (weight <= 0f) continue; + + fallbackTotalWeight += weight; + + if (Random.Shared.NextDouble() * fallbackTotalWeight < weight) + { + selected = item; + } } - if (arrayTotalWeight <= 0) + if (!hasElements) { - throw new InvalidOperationException("Total weight must be greater than zero."); + throw new InvalidOperationException("Sequence contains no elements."); } - var arrayTarget = (float)(Random.Shared.NextDouble() * arrayTotalWeight); - var arrayCumulative = 0f; - - for (int i = 0; i < arrayCount; i++) + if (fallbackTotalWeight <= 0f) { - var item = array[i]; - arrayCumulative += weightSelector(item); - if (arrayTarget <= arrayCumulative) - { - return item; - } + throw new InvalidOperationException("Total weight must be greater than zero."); } - return array[arrayCount - 1]; + return selected; } }