プログラミング
記事内に商品プロモーションを含む場合があります

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

Go言語で縦列列挙と組み合わせ列挙
Aru

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

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

Go言語でNextPermutation, Combinations

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

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

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]の2つしか出力されません

すべて列挙させる場合は、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にしています。ここは趣味です。

処理速度がより高速なコードはこちら

素直に組み合わせのパターンを列挙するコードを作成していたのですが、abc 328-E問題combinationsの処理が間に合いませんでした(TLEになりました)。

以下のサイトを参考にして、修正したコードがこちらです。

参考にしたサイト:Golang で順列・組み合わせ・重複順列・重複組み合わせの列挙

元のcombinations関数と同じ使い方ができるようにしましたので、そのまま入れ替えて利用することが可能です。

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
}

別バージョンも用意しました

第1引数がlistではなく、nCkのcomb(n,k)の形式で呼び出せるバージョンも作成しました。こちらの場合は、comp(6,3)といった形で呼び出せるので、リストを用意する必要はありません(ただし、戻り値は[0,1,2]といった0からn-1までの値の数列に固定されます)

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を1つずつ進める
		}
	}
	go func() {
		defer close(c)
		rec(0, 0)
	}()

	return
}

まとめ

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

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

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

ABOUT ME
ある/Aru
ある/Aru
IT&機械学習エンジニア/ファイナンシャルプランナー(CFP®)
専門分野は並列処理・画像処理・機械学習・ディープラーニング。プログラミング言語はC, C++, Go, Pythonを中心として色々利用。現在は、Kaggle, 競プロなどをしながら悠々自適に活動中
記事URLをコピーしました