Go言語で順列列挙と組み合わせ列挙を実装する（NextPermutation, Combinations)

Go言語には、C++のnext_permutationに相当する標準ライブラリが用意されていません。そこで、自分で、順列列挙（NextPermutation）と組み合わせ列挙（Combinations）を実装しました。この記事では、これらの機能を実現するソースコードと、その使い方を解説します。

その他のAtCoderに役立つ記事の一覧

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

Go言語でNextPermutation, Combinations

Go言語でAtCoderのコンテストに参加している人はまだ少なく、情報が不足している印象があります。私自身、最初のうちは海外サイトを含めて情報を探し、見つからない場合はC++やPythonのコードをGoに移植して使っていました。

Go言語を使う人が増えてほしいという思いから、私もAtCoderで使用している関数やライブラリの情報を発信していこうと思います。この記事では、順列列挙（NextPermutation）や組み合わせ列挙（Combinations）をGoで実装する方法を解説します。

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言語にはなかったので作成しましたが、テンプレとして活用していただいて構いません。

ある/Aru

IT＆機械学習エンジニア/ファイナンシャルプランナー（CFP®)

XContact

専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中