TypeScript で combinations を実装する

Article

#競プロ

#TypeScript

syakoo

Python の itertools.combinations のようなものを TypeScript で実装しました。

コード

function* combinations<T>(iterable: Iterable<T>, r: number) {
  const pool = [...iterable];
  const n = pool.length;
  if (n < r) return;
 
  const indices = [...new Array(r)].map((_, i) => i);
  yield indices.map((i) => pool[i]);
  while (true) {
    let i;
    for (i = r - 1; i >= 0; i--) {
      if (indices[i] !== n - (r - i)) {
        break;
      }
    }
    if (i === -1) return;
    indices[i]++;
    for (let j = i + 1; j < r; j++) {
      indices[j] = indices[j - 1] + 1;
    }
    yield indices.map((i) => pool[i]);
  }
}

配列やジェネレータなどの Iterable から r 個の組み合わせを取得するジェネレーターです。

使い方

0, 1, 2, ..., 4 から 3 個の組み合わせを持ってくる

const xs = [0, 1, 2, 3, 4];
const combs = [...combinations(xs, 3)];
/**
 * combs = [
 *   [0, 1, 2],
 *   [0, 1, 3],
 *   [0, 1, 4],
 *   [0, 2, 3],
 *   [0, 2, 4],
 *   [0, 3, 4],
 *   [1, 2, 3],
 *   [1, 2, 4],
 *   [1, 3, 4],
 *   [2, 3, 4],
 * ]
 **/

部分集合を取得する

const xs = ["A", "B", "C"];
const subsets = [[]];
 
xs.forEach((_, i) => subsets.push(...combinations(xs, i + 1)));
/**
 * subsets = [
 *   [],
 *   [ 'A' ],
 *   [ 'B' ],
 *   [ 'C' ],
 *   [ 'A', 'B' ],
 *   [ 'A', 'C' ],
 *   [ 'B', 'C' ],
 *   [ 'A', 'B', 'C' ]
 * ]
 **/