add a function to compute combinations with replacement

3 files changed, 106 insertions, 0 deletions

diff --git a/src/misc.c b/src/misc.c
index 183f0ce..88e4b77 100644
--- a/src/misc.c
+++ b/src/misc.c
@@ -684,3 +684,61 @@ void unique_vertices(int *id, size_t n, size_t npeaks, size_t *result, size_t *n
     free(ordered_results);
     free(unique_results);
 }
+
+static int is_sorted(size_t *a, size_t n)
+{
+    size_t i;
+
+    for (i = 1; i < n; i++)
+        if (a[i] < a[i-1]) return 0;
+
+    return 1;
+}
+
+void combinations_with_replacement(size_t n, size_t r, size_t *result, size_t *len)
+{
+    /* Returns the set of all unique combinations of length `r` from `n` unique
+     * elements.
+     *
+     * The result array can be indexed as:
+     *
+     *     result[i*r + j]
+     *
+     * where i is the ith combination and j is the jth element in the
+     * combination. For example:
+     *
+     *     size_t result[12];
+     *     size_t len;
+     *     product(3,2,result,&len);
+     *     for (i = 0; i < len; i++)
+     *         print("%zu %zu\n", result[i*2], result[i*2+1]);
+     *
+     * will print
+     *
+     *     0 0
+     *     0 1
+     *     0 2
+     *     1 1
+     *     1 2
+     *     2 2
+     *
+     * `result` should be an array with at least (n+r-1)!/r!/(n-1)! elements. */
+    size_t i, j;
+    size_t *tmp;
+
+    tmp = malloc(sizeof(size_t)*n*ipow(r,n));
+
+    product(n,r,tmp);
+
+    *len = 0;
+    for (i = 0; i < ipow(n,r); i++) {
+        if (is_sorted(tmp+i*r,r)) {
+            for (j = 0; j < n; j++) {
+                result[(*len)*r+j] = tmp[i*r+j];
+            }
+            *len += 1;
+        }
+    }
+
+    free(tmp);
+}
diff --git a/src/misc.h b/src/misc.h
index 8450454..02df398 100644
--- a/src/misc.h
+++ b/src/misc.h
@@ -30,5 +30,6 @@ double gamma_pdf(double x, double k, double theta);
 size_t ipow(size_t base, size_t exp);
 void product(size_t n, size_t r, size_t *result);
 void unique_vertices(int *id, size_t n, size_t npeaks, size_t *result, size_t *nvertices);
+void combinations_with_replacement(size_t n, size_t r, size_t *result, size_t *len);
 
 #endif
diff --git a/src/test.c b/src/test.c
index dcfc03c..85d266b 100644
--- a/src/test.c
+++ b/src/test.c
@@ -1838,6 +1838,52 @@ int test_find_peaks_sorted(char *err)
     return 0;
 }
 
+int test_combinations_with_replacement(char *err)
+{
+    /* Tests the combinations_with_replacement() function. */
+    size_t i, j;
+    size_t result[100];
+    size_t len;
+
+    size_t expected1[3][2] = {{0,0},{0,1},{1,1}};
+
+    combinations_with_replacement(2,2,result,&len);
+
+    if (len != 3) {
+        sprintf(err, "combinations_with_replacement() returned %zu combinations but expected %i", len, 3);
+        return 1;
+    }
+
+    for (i = 0; i < len; i++) {
+        for (j = 0; j < 2; j++) {
+            if (result[i*2 + j] != expected1[i][j]) {
+                sprintf(err, "result[%zu][%zu] = %zu but expected %zu", i, j, result[i*2+j], expected1[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    size_t expected2[6][2] = {{0,0},{0,1},{0,2},{1,1},{1,2},{2,2}};
+
+    combinations_with_replacement(3,2,result,&len);
+
+    if (len != 6) {
+        sprintf(err, "combinations_with_replacement() returned %zu combinations but expected %i", len, 6);
+        return 1;
+    }
+
+    for (i = 0; i < len; i++) {
+        for (j = 0; j < 2; j++) {
+            if (result[i*2 + j] != expected2[i][j]) {
+                sprintf(err, "result[%zu][%zu] = %zu but expected %zu", i, j, result[i*2+j], expected1[i][j]);
+                return 1;
+            }
+        }
+    }
+
+    return 0;
+}
+
 struct tests {
     testFunction *test;
     char *name;
@@ -1884,6 +1930,7 @@ struct tests {
     {test_unique_vertices, "test_unique_vertices"},
     {test_find_peaks_highest, "test_find_peaks_highest"},
     {test_find_peaks_sorted, "test_find_peaks_sorted"},
+    {test_combinations_with_replacement, "test_combinations_with_replacement"},
 };
 
 int main(int argc, char **argv)