Fixed GetRandomWeightedEntry throwing an exception when the random value is very close to 1

Author: Ellpeck

CHANGELOG.md

@@ -6,6 +6,7 @@ MLEM uses [semantic versioning](https://semver.org/). Potentially breaking chang
 ### MLEM
 Fixes
 - Fixed sprite animation groups not advancing to their first frame immediately
+- Fixed GetRandomWeightedEntry throwing an exception when the random value is very close to 1
 
 ### MLEM.Ui
 Additions

MLEM/MLEM.csproj

@@ -30,4 +30,8 @@
         <None Include="../Media/Logo.png" Pack="true" PackagePath="" />
         <None Include="../README.md" Pack="true" PackagePath="" />
     </ItemGroup>
+
+    <ItemGroup>
+        <InternalsVisibleTo Include="Tests" />
+    </ItemGroup>
 </Project>

MLEM/Maths/RandomExtensions.cs

@@ -8,6 +8,8 @@ namespace MLEM.Maths {
     /// </summary>
     public static class RandomExtensions {
 
+        private const string IndexOutOfRangeString = "Reached end of the collection in GetRandomWeightedEntry. For each entry, the passed weight function should return the same value on each invocation.";
+
         /// <summary>
         /// Gets a random entry from the given collection with uniform chance.
         /// </summary>
@@ -80,21 +82,21 @@ internal static T GetRandomWeightedEntry<T>(ICollection<T> entries, Func<T, int>
             var currWeight = 0;
             foreach (var entry in entries) {
                 currWeight += weightFunc(entry);
-                if (currWeight > goalWeight)
+                if (currWeight >= goalWeight)
                     return entry;
             }
-            throw new IndexOutOfRangeException();
+            throw new IndexOutOfRangeException(RandomExtensions.IndexOutOfRangeString);
         }
 
         internal static T GetRandomWeightedEntry<T>(ICollection<T> entries, Func<T, float> weightFunc, float randomValue) {
             var goalWeight = randomValue * entries.Sum(weightFunc);
             var currWeight = 0F;
             foreach (var entry in entries) {
                 currWeight += weightFunc(entry);
-                if (currWeight > goalWeight)
+                if (currWeight >= goalWeight)
                     return entry;
             }
-            throw new IndexOutOfRangeException();
+            throw new IndexOutOfRangeException(RandomExtensions.IndexOutOfRangeString);
         }
 
     }

Tests/CollectionTests.cs

@@ -1,3 +1,7 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using MLEM.Maths;
 using MLEM.Misc;
 using NUnit.Framework;
 
@@ -28,4 +32,46 @@ public void TestCombinations() {
         Assert.AreEqual(things.IndexCombinations(), indices);
     }
 
+    [Test]
+    public void TestRandomWeightedEntryEqual([Values(0.1F, 1, 20, 23.5F, 99, 10000000)] float equalWeight, [Values(true, false)] bool integer) {
+        var entries = new[] {"A", "B", "C", "D", "E"};
+        var random = new Random(390453);
+        var matches = new Dictionary<string, int>();
+        for (var i = 0; i < 100000; i++) {
+            var entry = integer ? random.GetRandomWeightedEntry(entries, _ => equalWeight.Ceil()) : random.GetRandomWeightedEntry(entries, _ => equalWeight);
+            matches[entry] = matches.GetValueOrDefault(entry) + 1;
+        }
+        for (var i = 0; i < entries.Length; i++)
+            Assert.AreEqual(100000 / entries.Length, matches[entries[i]], 1000);
+    }
+
+    [Test]
+    public void TestRandomWeightedEntryVaried([Values(1, 37.283923F, 99)] float weightMult, [Values(true, false)] bool integer) {
+        var weights = new[] {
+            ("A", 1),
+            ("B", 2),
+            ("C", 3),
+            ("D", integer ? 14 : 14.389238F),
+            ("E", 20)
+        };
+        var random = new Random(234598223);
+        var matches = new Dictionary<string, int>();
+        for (var i = 0; i < 1000000; i++) {
+            var entry = (integer ? random.GetRandomWeightedEntry(weights, e => (e.Item2 * weightMult).Ceil()) : random.GetRandomWeightedEntry(weights, e => e.Item2 * weightMult)).Item1;
+            matches[entry] = matches.GetValueOrDefault(entry) + 1;
+        }
+        for (var i = 0; i < weights.Length; i++) {
+            var expected = 1000000F / weights.Select(w => w.Item2).Sum() * weights[i].Item2;
+            Assert.AreEqual(expected, matches[weights[i].Item1], 1000);
+        }
+    }
+
+    [Test]
+    public void TestRandomWeightedEntryFixedValues() {
+        Assert.AreEqual(RandomExtensions.GetRandomWeightedEntry([1, 2, 3], _ => 1, 0), 1);
+        Assert.AreEqual(RandomExtensions.GetRandomWeightedEntry([1, 2, 3], _ => 1, 0.5F), 2);
+        Assert.AreEqual(RandomExtensions.GetRandomWeightedEntry([1, 2, 3], _ => 1, 0.99999999999999989F), 3);
+        Assert.AreEqual(RandomExtensions.GetRandomWeightedEntry([1, 2, 3], _ => 1, 1), 3);
+    }
+
 }