Go言語で順列列挙（NextPermutation）、組み合わせ列挙(Combinations)を実装

Go言語には、C++のnext_permutation, Pythonのitertool.permutationsに相当する、順列列挙の関数がありません。ということで自作してみましたので参考にしてください。

その他のAtCoderに役立つ記事の一覧は以下にあります。

あわせて読みたい

AtCoderで役立つ記事一覧（Python, Go言語）

Go言語でNextPermutation, Combinations

Go言語でAtCoderのコンテストに参加している人も少なく、情報が少なめかと思います。自分も最初のうちは、海外サイトを含めて情報を検索したり、それでも見つからない場合はC++やPythonのコードを移植したりしてきました。

Go言語を使う人が増えてほしいので、私の方でもAtCoderで使っている関数やライブラリなどの情報を発信していこうと思います。

この記事では、順列列挙、組み合わせ列挙を行う方法に関して説明します。

Go言語でAtCoderに参加する場合の一般的な注意点については以下の記事を参考にしてください。

あわせて読みたい

Go言語でAtCoderに参加する場合の注意点｜テンプレートも公開

NextPermutation

NextPermutation()は順列を列挙する関数です。

入力xに対して、次の順列を出力します。NextPermutation()がfalseを返したところで終了となります。

NextPermutations()の引数は、sort.Intaerfaceなので配列をsort.IntSlice()で変換して渡す必要があることに注意してください。

他に注意する点は、この関数はすべて列挙するものではなく、次を検索するものであるということです。例えば、x=[]int{4,3,1,2}からスタートさせると、[4,3,1,2], [4,3,2,1]の２つしか出力されません。

すべて列挙させる場合は、x = [0,1,2]やx=[1,2,3]といったように初期値が昇順に並んでいる必要があります。

func NextPermutation(x sort.Interface) bool {
	n := x.Len() - 1
	if n < 1 {
		return false
	}
	j := n - 1
	for ; !x.Less(j, j+1); j-- {
		if j == 0 {
			return false
		}
	}
	l := n
	for !x.Less(j, l) {
		l--
	}
	x.Swap(j, l)
	for k, l := j+1, n; k < l; {
		x.Swap(k, l)
		k++
		l--
	}
	return true
}

func main() {
	x := []int{1, 2, 3, 4}
	for {
		fmt.Println(x)
		if !NextPermutation(sort.IntSlice(x)) {
			break
		}
	}
}

出力結果

[1 2 3 4]
[1 2 4 3]
[1 3 2 4]
[1 3 4 2]
[1 4 2 3]
[1 4 3 2]
[2 1 3 4]
[2 1 4 3]
[2 3 1 4]
[2 3 4 1]
[2 4 1 3]
[2 4 3 1]
[3 1 2 4]
[3 1 4 2]
[3 2 1 4]
[3 2 4 1]
[3 4 1 2]
[3 4 2 1]
[4 1 2 3]
[4 1 3 2]
[4 2 1 3]
[4 2 3 1]
[4 3 1 2]
[4 3 2 1]

列挙するには昇順の必要があるので、ターゲットの変数aを順番に出力したい場合には、以下のような感じで書く必要があります。xをインデックスとして0〜len(a)-1の変数とし、xに従ってaを参照するという方法です。

func main() {
	a := []int{100, 3, 44, 23}
	x := []int{0, 1, 2, 3}
	for {
		for _, e := range x {
			fmt.Print(a[e], " ")
		}
		fmt.Println()

		if !NextPermutation(sort.IntSlice(x)) {
			break
		}
	}
}

Combinations

combinations(list []int, choose, buf int)関数は、組み合わせを列挙する関数です。

組み合わせ列挙に関しては、すべて列挙したリストを作成して戻すような方法を最初考えていました。ただ、これだとメモリを結構食うので、最終的に少し特殊な実装を行いました。

引数のlistは配列、chooseは取り出す数、bufは制御用のチャネルのサイズです。チャネルのサイズは1か2あたりを設定しておけば良いです。

$_nC_r$の$n$がlen(list)、$r$がchooseになります。

下のプログラムの例だと、$_5C_3=10個$を列挙しています。

こちらの関数は、列挙する関数なので、xの並びは特に意識する必要はありません。

//Combination generator for int slice
func combinations(list []int, choose, buf int) (c chan []int) {
	c = make(chan []int, buf)
	go func() {
		defer close(c)
		switch {
		case choose == 0:
			c <- []int{}
		case choose == len(list):
			c <- list
		case len(list) < choose:
			return
		default:
			for i := 0; i < len(list); i++ {
				for subComb := range combinations(list[i+1:], choose-1, buf) {
					c <- append([]int{list[i]}, subComb...)
				}
			}
		}
	}()
	return
}

func main() {
	x := []int{1, 2, 3, 4, 5}
	buf := 2

	fmt.Println("Combinations")
	for comb := range combinations(x, 3, buf) {
		fmt.Println(comb)
	}
}

出力結果：

[1 2 3]
[1 2 4]
[1 2 5]
[1 3 4]
[1 3 5]
[1 4 5]
[2 3 4]
[2 3 5]
[2 4 5]
[3 4 5]

実装に関して補足すると、combinations関数とfor文の間でチャネルを使った通信が行われています。バッファは1でいいのですが、conbinationsの中でgoを使ってスレッドを作って並列動作させているのでなんとなく2にしています。ここは趣味です。

func combinations(list []int, k, buf int) (c chan []int) {
	c = make(chan []int, buf)
	n := len(list)

	pattern := make([]int, k)

	var body func(pos, begin int)
	body = func(pos, begin int) {
		if pos == k {
			t := make([]int, k)
			copy(t, pattern)
			c <- t
			return
		}

		for num := begin; num < n+pos-k+1; num++ {
			pattern[pos] = list[num]
			body(pos+1, num+1)
		}
	}
	go func() {
		defer close(c)
		body(0, 0)
	}()

	return
}

func comb(n, k int) (c chan []int) {
	pat := make([]int, k)
	c = make(chan []int, 1)

	var rec func(pos, start int)

	rec = func(pos, start int) {
		// k個選んでいれば、それを出力する
		if pos == k {
			tmp := make([]int, k)
			copy(tmp, pat)
			c <- tmp
			return
		}
		// 選んでいない場合は、追加して再帰
		// 次に選べるのは、startからn-1までの値のいずれか
		for i := start; i < n; i++ {
			pat[pos] = i    // posに選んだ数字をセットして
			rec(pos+1, i+1) // pos, startを１つずつ進める
		}
	}
	go func() {
		defer close(c)
		rec(0, 0)
	}()

	return
}

まとめ

Go言語で、順列列挙と、組み合わせ列挙を行う関数を紹介しました。アルゴリズム的には、色々なところで説明されている通りの実装になります。

PythonやC++では用意されている関数ですが、Go言語にはなかったので作成しましたが、テンプレとして活用していただいて構いません。