diff --git a/src/index.ts b/src/index.ts index e5320f26..0cf6e32a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -23,6 +23,7 @@ export { compact, compactObject, first, flatmap, icompact } from "./custom"; export { chain, compress, + combinations, count, cycle, dropwhile, diff --git a/src/itertools.ts b/src/itertools.ts index a543b657..4daccb2f 100644 --- a/src/itertools.ts +++ b/src/itertools.ts @@ -35,6 +35,57 @@ export function compress(data: Iterable, selectors: Iterable): 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(iterable: Iterable, r?: number): IterableIterator { + 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 { }); }); +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]);