WeightedList update

Runtime/IsExternalInit.cs

@@ -0,0 +1,7 @@
+using System.ComponentModel;
+
+namespace System.Runtime.CompilerServices
+{
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    record IsExternalInit;
+}

Runtime/WeightedList/WeightedList.cs

@@ -4,34 +4,122 @@
 
 namespace TKLibs
 {
-  public static class WeightedList
+  public record WeightedItem<T>(T Item, int Weight);
+
+  public class WeightedList<T>
   {
-    public static T RandomWeightedItem<T>(this List<WeightedItem<T>> list, Random random = null)
+    readonly Dictionary<T, WeightedItem<T>> _itemCache = new();
+    readonly List<WeightedItem<T>> _itemList = new();
+    readonly Random _random;
+
+    int _currentMaxWeight = 0;
+
+    /// <summary>
+    /// Create a new WeightedList.
+    /// </summary>
+    /// <param name="random">Provide predefined random for procedural behavior. If null a default Random is used.</param>
+    public WeightedList(Random random = null)
+    {
+      _random = random ?? new Random();
+    }
+
+    /// <summary>
+    /// Get a random item from the list based on their weight.
+    /// An item with weight value of 2 will have double the chance to be chosen compared with an item with weight 1.
+    /// 
+    /// If removal of result item is desired, use RemoveRandomItem
+    /// </summary>
+    /// <returns>Random item based on weight</returns>
+    /// <exception cref="Exception"></exception>
+    public T RandomItem()
     {
-      random ??= new Random();
+      if(_currentMaxWeight <= 0) throw new IndexOutOfRangeException(
+        $"Unable to get random item in list because list is empty or all items are weightless. Total list weight is: {_currentMaxWeight}"
+      );
 
-      var ratio = list.Sum(i => i.Weight) / 100;
-
-      if (ratio == 0)
-        throw new Exception(
-          "Unable to find random item in list. Total list weight is: " + list.Sum(e => e.Weight)
-        );
-
-      var ran = random.Next(0, 100);
-      float weight = 0;
-
-      foreach (WeightedItem<T> item in list)
+      var ran = _random.Next(0, _currentMaxWeight);
+      var weight = 0;
+      foreach (var item in _itemList)
       {
-        weight += item.Weight / ratio;
-        if (ran < weight)
-        {
-          return item.Item;
-        }
+        weight += item.Weight;
+        if (ran < weight) return item.Item;
       }
-
+      
       throw new Exception(
-        "Unable to find random item in list. Total list weight is: " + list.Sum(e => e.Weight)
+        $"Unable to get random item in list. This is likely due to an error with the library. Random value is {ran}, total list weight is {_currentMaxWeight}"
       );
     }
+
+    /// <summary>
+    /// Remove a random item from the list based on their weight.
+    /// An item with weight value of 2 will have double the chance to be chosen compared with an item with weight 1.
+    ///
+    /// If removal of result item is NOT desired, use RandomItem
+    /// </summary>
+    /// <returns></returns>
+    public T RemoveRandomItem()
+    {
+      var item = RandomItem();
+      Remove(item);
+      return item;
+    }
+
+    /// <summary>
+    /// Add or Replace an item.
+    /// If the same Item already exists in list, its Weight will be replaced with the provided item's Weight 
+    /// </summary>
+    /// <param name="weightedItem">Item to add to list</param>
+    /// <returns>The provided weightedItem</returns>
+    public WeightedItem<T> AddOrReplace(WeightedItem<T> weightedItem)
+    {
+      if (_itemCache.TryGetValue(weightedItem.Item, out var oldItem))
+      {
+        _itemList.Remove(oldItem);
+      }
+
+      _itemCache[weightedItem.Item] = weightedItem;
+      _itemList.Add(weightedItem);
+      _currentMaxWeight = _itemList.Sum(i => i.Weight);
+
+      return weightedItem;
+    }
+
+    /// <summary>
+    /// Add or Replace an item with default Weight value of 1.
+    /// If the same Item already exists in list, its Weight will be set to 1.
+    /// </summary>
+    /// <param name="item">Item to add to list</param>
+    /// <returns>A WeightedItem created from provided Item with Weight 1</returns>
+    public WeightedItem<T> AddOrReplace(T item) => AddOrReplace(new WeightedItem<T>(item, 1));
+
+    /// <summary>
+    /// Remove an item from the list
+    /// </summary>
+    /// <param name="item">Item key to remove</param>
+    /// <returns>true if item successfully removed, false otherwise</returns>
+    public bool Remove(T item)
+    {
+      if (!_itemCache.TryGetValue(item, out var weightedItem)) return false;
+
+      var result = _itemCache.Remove(item) && _itemList.Remove(weightedItem);
+
+      if (result) _currentMaxWeight = _itemList.Sum(i => i.Weight);
+
+      return result;
+    }
+
+    /// <summary>
+    /// Remove an item from the list
+    /// </summary>
+    /// <param name="weightedItem">Item key to remove. Weight is ignored in removal</param>
+    /// <returns>true if item successfully removed, false otherwise</returns>
+    public bool Remove(WeightedItem<T> weightedItem) => Remove(weightedItem.Item);
+
+    /// <summary>
+    /// Get the WeightedItem corresponding to provided Item key 
+    /// </summary>
+    /// <param name="item">Item key to look for</param>
+    /// <returns>WeightedItem in list if exists, or null otherwise</returns>
+    public WeightedItem<T> GetWeightedItem(T item) =>  _itemCache.TryGetValue(item, out var weightedItem) ? weightedItem : null;
   }
 }

Tests/Editor/TestWeightedList.cs

@@ -0,0 +1,171 @@
+using NUnit.Framework;
+using TKLibs;
+using Random = System.Random;
+
+public class TestWeightedList
+{
+    [Test]
+    public void TestWeightedList_SimpleFlow()
+    {
+        var weightedList = new WeightedList<string>();
+        weightedList.AddOrReplace(TEST_STRING_1);
+
+        var result = weightedList.RandomItem();
+
+        Assert.AreEqual(TEST_STRING_1, result);
+    }
+
+    [Test]
+    public void TestWeightedList_ProceduralRandom()
+    {
+        Random testRandom = new(TEST_SEED);
+        var weightedList = new WeightedList<string>(testRandom);
+        weightedList.AddOrReplace(TEST_STRING_2);
+        weightedList.AddOrReplace(TEST_STRING_1);
+
+        var result = weightedList.RandomItem();
+
+        Assert.AreEqual(TEST_STRING_2, result);
+    }
+
+    [Test]
+    public void TestWeightedList_RoughlyEqualChance()
+    {
+        Random testRandom = new(TEST_SEED);
+        var weightedList = new WeightedList<string>(testRandom);
+        weightedList.AddOrReplace(TEST_STRING_1);
+        weightedList.AddOrReplace(TEST_STRING_2);
+        weightedList.AddOrReplace(TEST_STRING_3);
+        weightedList.AddOrReplace(TEST_STRING_4);
+
+        var results = new[]{0,0,0,0};
+        const int TEST_LOOP_COUNT = 100000; // weight 1 should appear ~ 25000 times
+
+        for (var i = 0; i < TEST_LOOP_COUNT; i++)
+        {
+            switch (weightedList.RandomItem())
+            {
+                case TEST_STRING_1: 
+                    results[0]++;
+                    break;
+                case TEST_STRING_2: 
+                    results[1]++;
+                    break;
+                case TEST_STRING_3: 
+                    results[2]++;
+                    break;
+                case TEST_STRING_4: 
+                    results[3]++;
+                    break;
+            }
+        }
+
+        Assert.AreEqual(new[]{25037,25213,24771,24979}, results);
+    }
+
+    [Test]
+    public void TestWeightedList_WeightBasedChance()
+    {
+        Random testRandom = new(TEST_SEED);
+        var weightedList = new WeightedList<string>(testRandom);
+        weightedList.AddOrReplace(new WeightedItem<string>(TEST_STRING_1, 1));
+        weightedList.AddOrReplace(new WeightedItem<string>(TEST_STRING_2, 2));
+        weightedList.AddOrReplace(new WeightedItem<string>(TEST_STRING_3, 3));
+
+        var results = new[]{0,0,0};
+        const int TEST_LOOP_COUNT = 120000; // weight 1 should appear ~ 20000 times
+
+        for (var i = 0; i < TEST_LOOP_COUNT; i++)
+        {
+            switch (weightedList.RandomItem())
+            {
+                case TEST_STRING_1: 
+                    results[0]++;
+                    break;
+                case TEST_STRING_2: 
+                    results[1]++;
+                    break;
+                case TEST_STRING_3: 
+                    results[2]++;
+                    break;
+            }
+        }
+
+        Assert.AreEqual(new[]{20122,40231,59647}, results);
+    }
+
+    [Test]
+    public void TestWeightedList_CanAddDefaultWeight()
+    {
+        var weightedList = new WeightedList<string>();
+
+        var weightedItem = weightedList.AddOrReplace(TEST_STRING_1);
+
+        Assert.AreEqual(1, weightedItem.Weight);
+    }
+
+    [Test]
+    public void TestWeightedList_CanAddDefinedWeight()
+    {
+        var weightedList = new WeightedList<string>();
+
+        var weightedItem = weightedList.AddOrReplace( new WeightedItem<string>(TEST_STRING_1, 12));
+
+        Assert.AreEqual(12, weightedItem.Weight);
+    }
+
+    [Test]
+    public void TestWeightedList_CanReplace()
+    {
+        var weightedList = new WeightedList<string>();
+        weightedList.AddOrReplace(TEST_STRING_1);
+        var itemToReplace = new WeightedItem<string>(TEST_STRING_1, 6);
+
+        weightedList.AddOrReplace(itemToReplace);
+
+        Assert.AreEqual(itemToReplace, weightedList.GetWeightedItem(TEST_STRING_1));
+    }
+
+    [Test]
+    public void TestWeightedList_CanRemove()
+    {
+        var weightedList = new WeightedList<string>();
+        weightedList.AddOrReplace(TEST_STRING_1);
+
+        weightedList.Remove(TEST_STRING_1);
+
+        Assert.IsNull(weightedList.GetWeightedItem(TEST_STRING_1));
+    }
+
+    [Test]
+    public void TestWeightedList_RemovingOnlyConsiderKey()
+    {
+        var weightedList = new WeightedList<string>();
+        weightedList.AddOrReplace(TEST_STRING_1);
+
+        weightedList.Remove( new WeightedItem<string>(TEST_STRING_1, 12));
+
+        Assert.IsNull(weightedList.GetWeightedItem(TEST_STRING_1));
+    }
+
+    [Test]
+    public void TestWeightedList_RemoveRandom()
+    {
+        var weightedList = new WeightedList<string>();
+        weightedList.AddOrReplace(TEST_STRING_1);
+        weightedList.AddOrReplace(TEST_STRING_2);
+        weightedList.AddOrReplace(TEST_STRING_3);
+        weightedList.AddOrReplace(TEST_STRING_4);
+
+        var removedItem = weightedList.RemoveRandomItem();
+
+        Assert.IsNull(weightedList.GetWeightedItem(removedItem));
+    }
+
+    const int TEST_SEED = 1337;
+
+    const string TEST_STRING_1 = "TEST_STRING_1";
+    const string TEST_STRING_2 = "TEST_STRING_2";
+    const string TEST_STRING_3 = "TEST_STRING_3";
+    const string TEST_STRING_4 = "TEST_STRING_4";
+}

Tests/Editor/avavt.tklibs-upm.Editor.Tests.asmdef

@@ -0,0 +1,24 @@
+{
+    "name": "avavt.tklibs-upm.Editor.Tests",
+    "rootNamespace": "",
+    "references": [
+        "UnityEngine.TestRunner",
+        "UnityEditor.TestRunner",
+        "avavt.tklibs-upm"
+    ],
+    "includePlatforms": [
+        "Editor"
+    ],
+    "excludePlatforms": [],
+    "allowUnsafeCode": false,
+    "overrideReferences": true,
+    "precompiledReferences": [
+        "nunit.framework.dll"
+    ],
+    "autoReferenced": false,
+    "defineConstraints": [
+        "UNITY_INCLUDE_TESTS"
+    ],
+    "versionDefines": [],
+    "noEngineReferences": false
+}