feat: add combinations

src/index.ts

@@ -23,6 +23,7 @@ export { compact, compactObject, first, flatmap, icompact } from "./custom";
 export {
   chain,
   compress,
+  combinations,
   count,
   cycle,
   dropwhile,

src/itertools.ts

@@ -35,6 +35,57 @@ export function compress<T>(data: Iterable<T>, selectors: Iterable<boolean>): T[
   return Array.from(icompress(data, selectors));
 }
 
+/**
+ * Return successive r-length combinations of elements in the iterable.
+ * 
+ * If `r` is not specified, then `r` defaults to the length of the iterable and
+ * all possible full-length combinations are generated.
+*/
+export function* combinations<T>(iterable: Iterable<T>, r?: number): IterableIterator<T[]>  {
+  const pool = Array.from(iterable);
+  const n = pool.length;
+  const x = r ?? n;
+
+  if (x < 0) {
+    throw Error("r must be non-negative");
+  }
+
+  if (x > n) {
+    return;
+  }
+
+  const indices: number[] = Array.from(range(n));
+  const result = Array(x);
+
+  let i: number;
+
+  while(true) {
+    for (i = 0; i < x; i++) {
+      const index = indices[i];
+      result[i] = pool[index];
+    }
+
+    yield result.slice(0, x);
+
+    for (i = x-1; i >= 0 && indices[i] == i+pool.length-x; i--);
+
+    if (i < 0) {
+      break;
+    }
+
+    indices[i]++;
+
+    for (let j = i+1; j<x; j++) {
+      indices[j] = indices[j-1] + 1;
+    }
+
+    for (; i < x; i++) {
+      const index = indices[i];
+      result[i] = pool[index];
+    }
+  }
+}
+
 /**
  * Returns an iterator producing elements from the iterable and saving a copy
  * of each.  When the iterable is exhausted, return elements from the saved

test/itertools.test.ts

@@ -4,6 +4,7 @@ import {
   all,
   chain,
   compress,
+  combinations,
   count,
   cycle,
   dropwhile,
@@ -64,6 +65,37 @@ describe("compress", () => {
   });
 });
 
+describe("combinations", () => {
+  it("combinations on empty list", () => {
+    expect(Array.from(combinations([]))).toEqual([[]]);
+  });
+
+  it("combinations of unique values", () => {
+    expect(Array.from(combinations([1, 2]))).toEqual([
+      [1, 2],
+    ]);
+
+    expect(Array.from(combinations([1, 2, 3]))).toEqual([
+      [1, 2, 3],
+    ]);
+
+    expect(Array.from(combinations([2, 2, 3]))).toEqual([
+      [2, 2, 3],
+    ]);
+  });
+
+  it("combinations with r param", () => {
+    // r too big
+    expect(Array.from(combinations([1, 2], 5))).toEqual([]);
+
+    // prettier-ignore
+    expect(Array.from(combinations(range(4), 2))).toEqual([
+            [0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3],
+        ]);
+  });
+
+});
+
 describe("count", () => {
   it("default counter", () => {
     expect(take(6, count())).toEqual([0, 1, 2, 3, 4, 5]);