C++ implementation of an algorithm proposed by Mike Earnest on math.stackexchange

earnestrank.cpp

// Given all possible combinations of K values, each no greater than M, and whose sum is exactly N
// (zeros and different orderings are accounted, that is 1 + 4 and 4 + 1 are different combinations),
// for the known combination, reconstruct back its index (the order of enumeration does not matter).
//
// C++ implementation of an algorithm proposed by Mike Earnest.
//
// Compile: g++-11 earnestrank.cpp -o earnestrank
// Run: ./earnestrank
//
//
// https://math.stackexchange.com/questions/4186930/get-the-index-of-combination-by-its-value
//
// There is a simple ranking algorithm with respect to the lexicographic ordering.
// Let _S_ be a sequence whose first entry is _i_, and where the deleting the initial _i_
// leaves the sequences _s'_ of length _k - 1_. Then all of the sequences coming before
// _S_ lexicographically fall in one of two categories:
//
// Sequences starting with _j_ where _j_ < _i_.
//
// Sequences starting with _i_, such that the substring _t_ left when you delete the initial
// _i_ occurs lexicographically before _s'_ (among sequences of length _k - 1_ summing to _n - i_).
//
// Letting _f(n, k, m)_ be the number of sequences of length _k_ with sum _m_ and maximum
// allowed entry _m_, it follows that rank(S) = rank(S') + Sum{j = 0, i -1} f(n - j, k - 1, m).

#include <array>
#include <cstdio>
#include <numeric>
#include <stdint.h>
#include <vector>

// Calculate the binomial coefficient "choose{n, k}"
// https://gist.github.com/shaunlebron/2843980
template <typename T>
T choose(uint32_t n, uint32_t k)
{
        if (0 == k || n == k) return 1;

        if (k > n) return 0;

        if (1 == k) return n;

        T b = 1;
        for (uint32_t i = 1; i <= k; ++i)
        {
                b *= (n - (k - i));
                b /= i;
        }

        return b;
}

// Find the number of representations of N as a sum of exactly K values not greater than M,
// including zeros and order-sensitive (1 + 4 and 4 + 1 are different representations).
// https://www.cyberforum.ru/post10391894.html
template <typename T>
T popcount(uint32_t n, uint32_t k, uint32_t m)
{
        int m1 = -1;
        T sum = 0;
        for (int i = 0, e = n / (m + 1); i <= e; i++)
        {
                m1 *= -1;
                sum += m1 * choose<T>(k, i) * choose<T>(n - i * (m + 1) + k - 1, k - 1);
        }

        return sum;
}

static uint32_t earnest_sequence_rank(uint32_t n, uint32_t k, uint32_t m, const uint32_t* sequence, size_t szsequence)
{
	// The rank of a single entry sequence is 1 (or zero, if you are zero-indexing).
	if (szsequence == 1) return 0;

	// Evaluate, each time omitting the leading term of sequence.
	uint32_t result = 0;
	uint32_t sum = std::accumulate(sequence, sequence + szsequence, 0);
	for (uint32_t i = 0; i < szsequence - 1; i++)
	{
		auto n_ = sum;
		auto k_ = szsequence - i;
		for (uint32_t j = 0; j < sequence[i]; j++)
			result += popcount<uint32_t>(n_ - j, k_ - 1, m);

		sum -= sequence[i];
	}
	return result;
}	

template<
	uint32_t n, // sequence elements sum
	uint32_t k, // length of sequence
	uint32_t m  // max allowed sequence element value
>
struct Combinations
{
        // Currying approach: https://stackoverflow.com/a/54508163/4063520
        template<uint32_t dim, class Callable>
        static constexpr void iterate(uint32_t sum, Callable&& c)
        {
                static_assert(dim > 0);
                for (uint32_t i = 0; i <= m; i++)
                {
                        if constexpr(dim == 1)
                        {
                                if (sum + i == n) c(i);
                        }
                        else
                        {
                                // Discarding overflowing paths.
                                if (sum + i <= n)
                                {
                                        auto bind_an_argument = [i, &c](auto... args)
                                        {
                                                c(i, args...);
                                        };

                                        iterate<dim - 1>(sum + i, bind_an_argument);
                                }
                        }
                }
        }

	template<class Callable>
	static void iterate(Callable&& c)
	{
		iterate<k>(0, c);
	}
};

int main(int argc, char* argv[])
{
	// Example setup.
	constexpr const uint32_t n = 16; // sequence elements sum
	constexpr const uint32_t k = 8; // length of sequence
	constexpr const uint32_t m = 10; // max allowed sequence element value

	uint64_t sum = 0;
	uint32_t min = static_cast<uint32_t>(-1), max = 0, count = 0;
	Combinations<n, k, m>::iterate([&](auto... args)
	{
		std::array<uint32_t, k> sequence = { args... };
		uint32_t rank = earnest_sequence_rank(n, k, m, sequence.data(), sequence.size());
		printf("rank = %u\n", rank);

		// Calculate statistics for correctness checking.
		min = std::min(min, rank);
		max = std::max(max, rank);
		sum += rank;
		count++;
	});

	if (min != 0)
	{
		fprintf(stderr, "Expected min = 0, but got %u\n", min);
		return -1;
	}
	if (max != count - 1)
	{
		fprintf(stderr, "Expected max = %u, but got %u\n", count - 1, max);
		return -1;
	}
	auto count_expected = popcount<uint32_t>(n, k, m);
	if (count != count_expected)
	{
		fprintf(stderr, "Expected count = %u, but got %u\n", count_expected, count);
		return -1;
	}
	auto sum_expected = (static_cast<uint64_t>(count) - 1) * count / 2;
	if (sum != sum_expected)
	{
		fprintf(stderr, "Expected sum = %u, but got %u\n", sum_expected, sum);
		return -1;
	}

	printf("OK!\n");
	return 0;
}