Skip to content

Instantly share code, notes, and snippets.

@cybercase
Last active February 10, 2023 10:59
Show Gist options
  • Save cybercase/db7dde901d7070c98c48 to your computer and use it in GitHub Desktop.
Save cybercase/db7dde901d7070c98c48 to your computer and use it in GitHub Desktop.
Python-like itertools.product function in javascript
function product() {
var args = Array.prototype.slice.call(arguments); // makes array from arguments
return args.reduce(function tl (accumulator, value) {
var tmp = [];
accumulator.forEach(function (a0) {
value.forEach(function (a1) {
tmp.push(a0.concat(a1));
});
});
return tmp;
}, [[]]);
}
console.log(product([1], [2, 3], ['a', 'b']));
@iota-pi
Copy link

iota-pi commented Jun 27, 2019

Generator version (including TypeScript definition), sans repeats

function* product<T> (...pools: T[][]): IterableIterator<T[]> {
// function* product (...pools) {
  let i = 0;
  const indexes = new Array(pools.length).fill(0)
  const result = indexes.map((x, i) => pools[i][x]);
  indexes[0] = -1;
  while (i < indexes.length) {
    if (indexes[i] < pools[i].length - 1) {
      indexes[i]++;
      result[i] = pools[i][indexes[i]];
      i = 0;

      // NB: can remove `.slice()` for a performance improvement if you don't mind mutating the same result object
      yield result.slice();
    } else {
      indexes[i] = 0;
      result[i] = pools[i][0];
      i++;
    }
  }
}

@colin-daniels
Copy link

Generator version for iterables, with tuple type output:

// NB: throws if any iterables passed to it are empty
function* product<T extends Array<Iterable<any>>>(...iterables: T): IterableIterator<{
    [K in keyof T]: T[K] extends Iterable<infer U> ? U : never
}> {
    if (iterables.length === 0) { return; }
    // make a list of iterators from the iterables
    const iterators = iterables.map(it => it[Symbol.iterator]());
    const results = iterators.map(it => it.next());
    if (results.some(r => r.done)) {
        throw new Error("Input contains an empty iterator.");
    }
    
    for (let i = 0;;) {
        if (results[i].done) {
            // reset the current iterator
            iterators[i] = iterables[i][Symbol.iterator]();
            results[i] = iterators[i].next();
            // advance, and exit if we've reached the end
            if (++i >= iterators.length) { return; }
        } else {
            yield results.map(({ value }) => value) as any;
            i = 0;
        }
        results[i] = iterators[i].next();
    }
}

let p: IterableIterator<[string, number, boolean]> = product(
    ['a', 'b'],
    [1, 2, 3],
    [true, false]
);

// [ 'a', 1, true  ]
// [ 'b', 1, true  ]
// [ 'a', 2, true  ]
// [ 'b', 2, true  ]
// [ 'a', 3, true  ]
// [ 'b', 3, true  ]
// [ 'a', 1, false ]
// [ 'b', 1, false ]
// [ 'a', 2, false ]
// [ 'b', 2, false ]
// [ 'a', 3, false ]
// [ 'b', 3, false ]
for (let v of p) { console.log(v); }

// [ 'a', 1 ]
// [ 'a', 2 ]
for (let v of product(['a'], [1, 2])) { console.log(v); }

// [ 'a', 1, 2 ]
// [ 'b', 1, 2 ]
// [ 'c', 1, 2 ]
for (let v of product(['a', 'b', 'c'], [1], [2])) { console.log(v); }

// No output
for (let v of product()) { console.log(v); }

@graup
Copy link

graup commented Apr 24, 2021

Same as above but with inferred output types

type Iterableify<T> = { [K in keyof T]: Iterable<T[K]> }

function* product<T extends Array<unknown>>(...iterables: Iterableify<T>): Generator<T> {
  if (iterables.length === 0) { return; }
  const iterators = iterables.map(it => it[Symbol.iterator]());
  const results = iterators.map(it => it.next());
  
  // Cycle through iterators
  for (let i = 0;;) {
    if (results[i].done) {
      // Reset the current iterator
      iterators[i] = iterables[i][Symbol.iterator]();
      results[i] = iterators[i].next();
      // Advance and exit if we've reached the end
      if (++i >= iterators.length) { return; }
    } else {
      yield results.map(({ value }) => value) as T;
      i = 0;
    }
    results[i] = iterators[i].next();
  }
}

@sponrad
Copy link

sponrad commented Dec 29, 2021

This is fantastic... it took me a little while to understand it so I rewrote the original with descriptive variable names and updated some syntax for ES6.

function arrayProduct(...arrays) {
    return arrays.reduce((prevAccumulator, currentArray) => {
        let newAccumulator = [];
        prevAccumulator.forEach(prevAccumulatorArray => {
            currentArray.forEach(currentValue => {
                newAccumulator.push(prevAccumulatorArray.concat(currentValue));
            });
        });
        return newAccumulator;
    }, [[]]);
}

@Hoppingmad9
Copy link

Doing arrays[0] (for @sponrad 's answer) or args[0] (for @cybercase 's) allows you to use Object.values(x) as the input. Useful if you don't know how many arrays are in your object/array.

function arrayProduct(...arrays) {
    return arrays[0].reduce((prevAccumulator, currentArray) => {
        let newAccumulator = [];
        prevAccumulator.forEach(prevAccumulatorArray => {
            currentArray.forEach(currentValue => {
                newAccumulator.push(prevAccumulatorArray.concat(currentValue));
            });
        });
        return newAccumulator;
    }, [[]]);
}

const packet = {
    'weight': ['w1', 'w2', 'w3', 'w4'],
    'speed': ['s1', 's2', 's3'],
    'colour': ['c1', 'c2'],
}

console.log(arrayProduct(Object.values(packet)));
// [["w1", "s1", "c1"], ["w1", "s1", "c2"], ["w1", "s2", "c1"], ["w1", "s2", "c2"], ["w1", "s3", "c1"], ["w1", "s3", "c2"], 
// ["w2", "s1", "c1"], ["w2", "s1", "c2"], ["w2", "s2", "c1"], ["w2", "s2", "c2"], ["w2", "s3", "c1"], ["w2", "s3", "c2"], 
// ["w3", "s1", "c1"], ["w3", "s1", "c2"], ["w3", "s2", "c1"], ["w3", "s2", "c2"],  ["w3", "s3", "c1"], ["w3", "s3", "c2"], 
// ["w4", "s1", "c1"], ["w4", "s1", "c2"], ["w4", "s2", "c1"], ["w4", "s2", "c2"], ["w4", "s3", "c1"], ["w4", "s3", "c2"]]

@gaspardfeuvray
Copy link

gaspardfeuvray commented Jun 11, 2022

Generic version for Typescript:

const combinations = <T>(sets: T[][]): T[][] => {
  if (sets.length === 1) {
    return sets[0].map((el) => [el]);
  } else
    return sets[0].flatMap((val) =>
      combinations(sets.slice(1)).map((c): T[] => [val].concat(c))
    );
};

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment