/* Enumerating k-weighted bitstrings of length n
 * Daniel Beer <dlbeer@gmail.com>, 4 Nov 2015
 *
 * This program is in the public domain.
 *
 * Note the following:
 *
 *     (n, k) = n*k     * (n-1, k-1)
 *     (n, k) = n*(n-k) * (n-1, k)
 *
 * Where (n, k) = n! / [(n-k)! * k!]
 */

#include <stdio.h>
#include <assert.h>

/* Return the weight of the bit-string s */
static int nbits(int s)
{
	int r = 0;

	while (s) {
		s &= s - 1;
		r++;
	}

	return r;
}

/* Compute (n, k) = n! / [(n-k)! * k!] */
static int nck(int n, int k)
{
	/* Start by computing (n-k, 0) = 1 */
	int r = 1;
	int i;

	/* At each step, compute (n-k+i, i) from (n-k+i-1, i-1) */
	for (i = 1; i <= k; i++)
		r = r * (i + n - k) / i;

	return r;
}

/* Given a bit-string s of length n and weight k, return its index in
 * the sorted list of all k-weighted bitstrings of length n.
 */
static int kw2i(int n, int k, int s)
{
	/* c is the number of items in this subtree: the subtree
	 * containing all k-weighted bitstrings of length n. Both n and
	 * k will change as we descend the tree.
	 */
	int c = nck(n, k);
	int r = 0;

	while (n) {
		/* Count the number of items in the left (0) and right
		 * (1) subtrees.
		 */
		const int c0 = c * (n - k) / n;
		const int c1 = c * k / n;

		/* Based on the bit in position n, descend into one of
		 * the subtrees. If it's the right subtree, add the size
		 * of the left subtree to the count of all bitstrings
		 * which are lexicographically before our given
		 * bitstring.
		 */
		if (s & (1 << (n - 1))) {
			r += c0;
			c = c1;
			k--;
		} else {
			c = c0;
		}

		n--;
	}

	return r;
}

/* Given an index into the list of all k-weighted bitstrings of length
 * n, return the rth bitstring.
 */
static int i2kw(int n, int k, int r)
{
	int c = nck(n, k);
	int s = 0;

	while (n) {
		const int c0 = c * (n - k) / n;
		const int c1 = c * k / n;

		/* This traces the same path through the tree of bits as kw2i,
		 * but it uses the index to decide which subtree to descend
		 * into. If the index is >= the number of items in the
		 * left subtree, we must descend into the right.
		 */
		if (r >= c0) {
			s |= 1 << (n - 1);
			r -= c0;
			c = c1;
			k--;
		} else {
			c = c0;
		}

		n--;
	}

	return s;
}

/* Test: find all k-weighted bitstrings of length n, in order, by brute
 * force. Check that they are assigned sequential indices, and that
 * those indices can retrieve the original bitstring.
 */
int main(void)
{
	const int N = 11;
	const int K = 4;
	int count = 0;
	int i;

	for (i = 0; i < (1 << N); i++) {
		if (nbits(i) == K) {
			const int r = kw2i(N, K, i);
			const int s = i2kw(N, K, r);

			printf("%04x ==> %3d ==> %04x\n", i, r, s);

			assert(count == r);
			assert(s == i);
			count++;
		}
	}

	assert(count == nck(N, K));
	return 0;
}