GoでC++のnext_permutationを実装するメモ

C++の next_permutation の動作を確認する

C++
#include <algorithm>
#include <vector>
#include <iostream>

std::vector<int> v = {1, 2, 3};

void print() {
  for (int &e: v) std::cout << e << " ";
  std::cout << std::endl;
}

int main() {
  do {
    print();
  } while (std::next_permutation(v.begin(), v.end()));

  std::cout << "最後: ";
  print();
}

出力

1 2 3
1 3 2
2 1 3
2 3 1
3 1 2
3 2 1
最後: 1 2 3

これを Go で実装する。

注意点は

  1. 全ての順列を生成し終えたあとは、初期の並びに戻っていること
  2. do-while でループしていること

Go で実装する

簡単な処理の流れメモ

  1. スライスの末尾から隣接する要素を見ていき、s[i] < s[i+1]となる最初の i を決定する
  2. スライスの末尾から要素を見ていき、s[i] < s[j]となる最初の j を決定する
  3. s[i]とs[j]を swap する
  4. s[i+1:]を reverse する

実装例

予め go get golang.org/x/exp で constraints.Ordered を使えるようにしている。

Go
package lib

import "golang.org/x/exp/constraints"

func Reverse[S ~[]E, E any](s S) {
	for i, j := 0, len(s)-1; i < j; i, j = i+1, j-1 {
		s[i], s[j] = s[j], s[i]
	}
}

func NextPermutation[S ~[]E, E constraints.Ordered](s S) bool {
	i := len(s) - 2
	for i >= 0 && s[i] >= s[i+1] {
		i--
	}

	if i < 0 {
		Reverse(s)
		return false
	}

	j := len(s) - 1
	for j > i && s[j] <= s[i] {
		j--
	}

	s[i], s[j] = s[j], s[i]
	Reverse(s[i+1:])

	return true
}
Go
package main

import (
	"fmt"
	"github.com/murosan/lib"
)

func main() {
	v := []int{1, 2, 3}

	for {
		fmt.Println(v)

		// do-while
		if !lib.NextPermutation(v) {
			break
		}
	}

	fmt.Println("最後:", v)
}

出力

[1 2 3]
[1 3 2]
[2 1 3]
[2 3 1]
[3 1 2]
[3 2 1]
最後: [1 2 3]

do-while の部分はこういう書き方でも良い。

Go
for do := true; do; do = lib.NextPermutation(s) {
  // 処理
}

5要素の中から2要素を選ぶパターンの全列挙、みたいなこともできる

Go
func main() {
	v := []int{0, 0, 0, 1, 1}

	for {
		fmt.Println(v)
		if !lib.NextPermutation(v) {
			break
		}
	}

	fmt.Println("最後:", v)
}

[0 0 0 1 1]
[0 0 1 0 1]
[0 0 1 1 0]
[0 1 0 0 1]
[0 1 0 1 0]
[0 1 1 0 0]
[1 0 0 0 1]
[1 0 0 1 0]
[1 0 1 0 0]
[1 1 0 0 0]
最後: [0 0 0 1 1]

テスト

Go
import (
	"testing"
	"golang.org/x/exp/slices"
)

func TestNextPermutation(t *testing.T) {
	cases := []struct {
		in []int
		ex []int
		b  bool
	}{
		{[]int{1, 2, 3}, []int{1, 3, 2}, true},
		{[]int{3, 2, 1}, []int{1, 2, 3}, false},
		{[]int{1, 1, 5}, []int{1, 5, 1}, true},
		{[]int{1}, []int{1}, false},
		{[]int{}, []int{}, false},
		{[]int{1, 3, 2}, []int{2, 1, 3}, true},
		{[]int{2, 3, 1}, []int{3, 1, 2}, true},
		{[]int{1, 5, 1}, []int{5, 1, 1}, true},
		{[]int{5, 1, 1}, []int{1, 1, 5}, false},
		{[]int{2, 1, 1, 4, 3}, []int{2, 1, 3, 1, 4}, true},
	}

	for i, c := range cases {
		result := NextPermutation(c.in)
		if result != c.b || !slices.Equal(c.in, c.ex) {
			t.Errorf("[%d] want: %v, actual: %v, bool:%v", i, c.ex, c.in, result)
		}
	}
}

おまけ

PrevPermutationの実装

Go
func PrevPermutation[S ~[]E, E constraints.Ordered](s S) bool {
	i := len(s) - 2
	for i >= 0 && s[i] <= s[i+1] {
		i--
	}

	if i < 0 {
		Reverse(s)
		return false
	}

	j := len(s) - 1
	for j > i && s[j] >= s[i] {
		j--
	}

	s[i], s[j] = s[j], s[i]
	Reverse(s[i+1:])

	return true
}

テスト

Go
func TestPrevPermutation(t *testing.T) {
	cases := []struct {
		in []int
		ex []int
		b  bool
	}{
		{[]int{1, 3, 2}, []int{1, 2, 3}, true},
		{[]int{1, 2, 3}, []int{3, 2, 1}, false},
		{[]int{1, 5, 1}, []int{1, 1, 5}, true},
		{[]int{1}, []int{1}, false},
		{[]int{}, []int{}, false},
		{[]int{2, 1, 3}, []int{1, 3, 2}, true},
		{[]int{3, 1, 2}, []int{2, 3, 1}, true},
		{[]int{5, 1, 1}, []int{1, 5, 1}, true},
		{[]int{1, 1, 5}, []int{5, 1, 1}, false},
		{[]int{2, 1, 3, 1, 4}, []int{2, 1, 1, 4, 3}, true},
	}

	for i, c := range cases {
		result := PrevPermutation(c.in)
		if result != c.b || !slices.Equal(c.in, c.ex) {
			t.Errorf("[%d] want: %v, actual: %v, bool:%v", i, c.ex, c.in, result)
		}
	}
}

もっと抽象化してみる

Go
func NextPermutationFunc[S ~[]E, E any](s S, comp func(a, b E) int) bool {
	i := len(s) - 2
	for i >= 0 && comp(s[i], s[i+1]) >= 0 {
		i--
	}

	if i < 0 {
		Reverse(s)
		return false
	}

	j := len(s) - 1
	for j > i && comp(s[i], s[j]) >= 0 {
		j--
	}

	s[i], s[j] = s[j], s[i]
	Reverse(s[i+1:])

	return true
}

こう実装すると、さっきは[]intで5要素の中から2要素を選ぶパターンを全列挙をしたが、[]boolでもできるようになる。

Go
func main() {
	v := []bool{false, false, false, true, true}

	comp := func(a, b bool) int {
		if a == b {
			return 0
		}
		if a {
			return 1
		}
		return -1
	}

	for {
		fmt.Println(v)
		if !lib.NextPermutationFunc(v, comp) {
			break
		}
	}

	fmt.Println("最後:", v)
}

[false false false true true]
[false false true false true]
[false false true true false]
[false true false false true]
[false true false true false]
[false true true false false]
[true false false false true]
[true false false true false]
[true false true false false]
[true true false false false]
最後: [false false false true true]