repeated_combination

Ruby の Array#repeated_combination は重複を許す組み合わせを返すメソッドだ(厳密にいうとEnumerator を返す)。
次のブログ記事で見つけた。

cf. 続:Haskell/Clojureでrepeated_combinationを実装してみる – Programmer’s Note

これを Go でやってみよう、ということで Haskell のコードを参考に作ってみた。

package main

import (
    "fmt"
)

func main() {
    xs := []int{ 1,2,3 }
    fmt.Println(repeatedCombination(xs, 3))
}

func repeatedCombination(xs []int, n int) [][]int {
    c := make([][]int, 0)
    if n == 1 {
        for _, x := range xs {
            c = append(c, []int{ x })
        }
    } else {
        for _, x := range xs {
            for _, y := range repeatedCombination(xs, n - 1) {
                c = append(c, append([]int{ x }, y...))
            }
        }
    }
    return c
}

ところが、実行結果が Array#repeated_combination と合わない。
上のコードの実行結果はこう。

^o^ > go run repeated_combination.go
[[1 1 1] [1 1 2] [1 1 3] [1 2 1] [1 2 2] [1 2 3] [1 3 1] [1 3 2] [1 3 3] [2 1 1] [2 1 2] [2 1 3] [2 2 1] [2 2 2] [2 2 3] [2 3 1] [2 3 2] [2 3 3] [3 1 1] [3 1 2] [3 1 3] [3 2 1] [3 2 2] [3 2 3] [3 3 1] [3 3 2] [3 3 3]]

対して Array#repeated_combination はこう。

irb(main):001:0> [1,2,3].repeated_combination(3).to_a
=> [[1, 1, 1], [1, 1, 2], [1, 1, 3], [1, 2, 2], [1, 2, 3], [1, 3, 3], [2, 2, 2], [2, 2, 3], [2, 3, 3], [3, 3, 3]]

順列ではなく組み合わせなので、Array#repeated_combination では例えば [1,1,2] と [1,2,1] と [2,1,1] を同じものとしているのに対して、Go のコードでは別のものとして数え上げてしまっている。
そこで改良を加えたのがこれだ。

package main

import (
    "fmt"
)

func main() {
    xs := []int{ 1,2,3 }
    fmt.Println(repeatedCombination(xs, 3))
}

func repeatedCombination(xs []int, n int) [][]int {
    c := make([][]int, 0)
    if n == 1 {
        for _, x := range xs {
            c = append(c, []int{ x })
        }
    } else {
        for _, x := range xs {
            for _, y := range repeatedCombination(xs, n - 1) {
                if x <= y[0] {
                    c = append(c, append([]int{ x }, y...))
                }
            }
        }
    }
    return c
}

実行結果。

^o^ > go run repeated_combination2.go
[[1 1 1] [1 1 2] [1 1 3] [1 2 2] [1 2 3] [1 3 3] [2 2 2] [2 2 3] [2 3 3] [3 3 3]]

今度はうまくいった。