From e4a1ede2d2dffaa36060b4e588e658fd063d1d5c Mon Sep 17 00:00:00 2001
From: lisyarus <lisyarus@gmail.com>
Date: Sat, 7 May 2022 18:44:30 +0300
Subject: [PATCH] Add random weighted distribution

---
 .../random/include/psemek/random/weighted.hpp | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 libs/random/include/psemek/random/weighted.hpp

diff --git a/libs/random/include/psemek/random/weighted.hpp b/libs/random/include/psemek/random/weighted.hpp
new file mode 100644
index 00000000..3b65a1a0
--- /dev/null
+++ b/libs/random/include/psemek/random/weighted.hpp
@@ -0,0 +1,43 @@
+#pragma once
+
+#include <psemek/random/uniform.hpp>
+
+#include <vector>
+#include <algorithm>
+
+namespace psemek::random
+{
+
+	template <typename T>
+	struct weighted_distribution
+	{
+		weighted_distribution(std::vector<T> frequencies)
+			: prefix_sum_(std::move(frequencies))
+		{
+			T sum = T{0};
+			for (auto const & w : prefix_sum_)
+				sum += w;
+
+			for (auto & w : prefix_sum_)
+				w /= sum;
+
+			for (std::size_t i = 1; i < prefix_sum_.size(); ++i)
+				prefix_sum_[i] += prefix_sum_[i - 1];
+		}
+
+		std::size_t size() const { return prefix_sum_.size(); }
+
+		template <typename RNG>
+		std::size_t operator()(RNG && rng) const
+		{
+			auto roll = uniform<T>(rng);
+			auto it = std::upper_bound(prefix_sum_.begin(), prefix_sum_.end(), roll);
+			if (it == prefix_sum_.end()) --it;
+			return it - prefix_sum_.begin();
+		}
+
+	private:
+		std::vector<T> prefix_sum_;
+	};
+
+}