WeightedRandom.cs

WeightedRandom.cs

using System;
using System.Collections;
using System.Collections.Generic;
using Random = UnityEngine.Random;

// ReSharper disable UnusedMember.Global
// ReSharper disable MemberCanBePrivate.Global
namespace Utils
{
	public class WeightedRandomizer<T> : IEnumerable
	{
		private readonly List<Entry> _pool = new List<Entry>();

		public WeightedRandomizer()
		{
		}

		public WeightedRandomizer(double weight, T item)
		{
			Add(weight, item);
		}

		public WeightedRandomizer(Entry entry)
		{
			Add(entry);
		}

		public T GetNext()
		{
			if (_pool.Count == 0) throw new ArgumentOutOfRangeException("Cannot call next on a empty pool !");
			double randomDice = Random.Range(0.0f, GetWeightsSum());
			foreach (Entry entry in _pool)
			{
				if (randomDice < entry.Weight)
				{
					return entry.Item;
				}
				randomDice -= entry.Weight;
			}
			throw new ArgumentOutOfRangeException("This code is unreachable");
		}

		private float GetWeightsSum()
		{
			double sum = 0;
			foreach (Entry entry in _pool)
			{
				sum += entry.Weight;
			}
			return (float) sum;
		}

		public void Add(double weight, T item)
		{
			_pool.Add(new Entry(weight, item));
		}

		public void Add(Entry entry)
		{
			_pool.Add(entry);
		}


		public class Entry
		{
			public double Weight { get; private set; }
			public T Item { get; private set; }

			public Entry(double weight, T item)
			{
				Weight = weight;
				Item = item;
			}
		}

		public IEnumerator GetEnumerator()
		{
			foreach (Entry entry in _pool)
			{
				yield return entry;
			}
		}
	}
}