Generalize Partition class, add unit tests

Author: danngreen

tests/partition_tests.cc

@@ -0,0 +1,21 @@
+#include "doctest.h"
+#include "util/partition.hh"
+#include <array>
+
+TEST_CASE("Basic usage") {
+
+	Partition<2, 8> parts;
+
+	auto nums = std::array{1u, 9u, 3u, 5u, 2u, 10u};
+	parts.calculate(nums);
+
+	CHECK(parts.parts[0].size() == 3);
+	CHECK(nums[parts.parts[0][0]] == 10);
+	CHECK(nums[parts.parts[0][1]] == 3);
+	CHECK(nums[parts.parts[0][2]] == 2);
+
+	CHECK(parts.parts[1].size() == 3);
+	CHECK(nums[parts.parts[1][0]] == 9);
+	CHECK(nums[parts.parts[1][1]] == 5);
+	CHECK(nums[parts.parts[1][2]] == 1);
+}

util/partition.hh

@@ -2,15 +2,25 @@
 #include "util/fixed_vector.hh"
 #include <span>
 
-template<unsigned MaxElements>
-struct PartitionTwo {
+// Splits a list of elements into N groups, the sum of each group being roughly equal.
+// The **indices** of the original elements are stored resulting `parts`
+template<unsigned NumPartitions, unsigned MaxElements, typename T = unsigned>
+struct Partition {
+	// parts[N] contains the indices of the elements whose sums are roughly equal to other parts[]
+	std::array<FixedVector<unsigned, MaxElements>, NumPartitions> parts{};
 
-	FixedVector<unsigned, MaxElements> a{};
-	FixedVector<unsigned, MaxElements> b{};
+	Partition() = default;
+
+	Partition(std::span<T> vals) {
+		calculate(vals);
+	}
+
+	void calculate(std::span<T> vals) {
+		for (auto &part : parts)
+			part.clear();
 
-	PartitionTwo(std::span<unsigned> vals) {
 		struct IdVal {
-			unsigned val;
+			T val;
 			unsigned id;
 		};
 
@@ -21,18 +31,14 @@ struct PartitionTwo {
 
 		std::ranges::sort(ordered, std::greater{}, &IdVal::val);
 
-		unsigned a_sum = 0;
-		unsigned b_sum = 0;
+		std::array<T, NumPartitions> sums{};
 
-		// The part with the smaller sum gets the next element
+		// The part with the smallest sum gets the next element
 		for (auto [val, i] : ordered) {
-			if (a_sum <= b_sum) {
-				a_sum += val;
-				a.push_back(i);
-			} else {
-				b_sum += val;
-				b.push_back(i);
-			}
+			auto min_sum = std::ranges::min_element(sums);
+			auto idx = std::ranges::distance(sums.begin(), min_sum);
+			*min_sum += val;
+			parts[idx].push_back(i);
 		}
 	}
 };