順列を全列挙する関数を実装する(Swift)

最初に

競プロの問題を解いていると考えられる順列を全列挙したい時がある。毎回実装すると時間を食ってしまうので関数にして時短したい。

考え方

はじめに[1, 2]の順列を全列挙する場合を考えてみる。考えられる順列は、

1 – 2
2 – 1

の2通りである。

次に[0, 1, 2]の順列を全列挙する場合を考える。

0 – 1 – 2
0 – 2 – 1
2 – 0 – 1
2 – 1 – 0
1 – 0 – 2
1 – 2 – 0

の6通りである。

ここで0 – 1 – 2、 0 – 2 – 1をみると、[1, 2]の順列を全列挙したものの先頭に0がくっついていることがわかる。0 – ([1, 2]の順列全列挙)というわけである。

残りの 2 – 0 – 1、2 – 1 – 0、1 – 0 – 2、1 – 2 – 0も同様に2 – ([0, 1]の順列全列挙)、 1 – ([0, 2]の順列全列挙)で求めることができる。このことから、先頭の値 – (残りの値の順列全列挙)ですべての順列を求めることが可能だとわかる。

つまり、順列を全列挙するためには順列を全列挙する必要があるということである。再帰関数で実装すれば求められそうである。

コードに落とし込む

最初に、完成したコードを示す。後に各行で何を行っているのかを示す。

func permutation(_ args: [Int]) -> [[Int]] {
  guard args.count > 1 else { return [args] }
  func rotate(_ arr: [Int]) -> [Int] { [arr.last!] + arr.dropLast() }
  var rotatedValue = args
  var result: [[Int]] = []
    for _ in 0..<args.count {
      let head = rotatedValue.first!
      let tail = rotatedValue.dropFirst().map { Int($0) }
      permutation(tail).forEach {
        result.append([head] + $0)
      }
      rotatedValue = rotate(rotatedValue)
    }
    return result
}
print(permutation([0, 1, 2]))  //[[0, 1, 2], [0, 2, 1], [2, 0, 1], [2, 1, 0], [1, 2, 0], [1, 0, 2]]

2行目

guard args.count > 1 else { return [args] }

引数のargsで与えられた配列の要素数が0または1の場合、それが順列を全列挙したものなのでそのまま二次元配列にして返す。

3行目

func rotate(_ arr: [Int]) -> [Int] { [arr.last!] + arr.dropLast() }

考え方でも述べた通り0 – ([1, 2]の順列全列挙)、2 – ([0, 1]の順列全列挙)…というふうに順列を全列挙させる。そのため、先頭の値を変える必要がある。この関数は、引数で与えられた配列の最後の要素を先頭に移動させて配列を返している。

4行目

var rotatedValue = args

先頭の値を変えた配列を入れる変数である。argsで初期化している。

6行目

for _ in 0..<args.count {

与えられた配列が[x, y, z]の場合は先頭の値が、xのとき、yのとき、zのときの計3回を考える必要がある。この回数は、argsの要素数と一致する。

7行目、8行目

let head = rotatedValue.first!
let tail = rotatedValue.dropFirst().map { Int($0) }

先頭の値と残りの値で分ける。

9行目、10行目、11行目

permutation(tail).forEach {
  result.append([head] + $0)
}

残りの値の順列全列挙を求め、先頭の値にくっつけてresultへ追加する。

12行目

rotatedValue = rotate(rotatedValue)

先頭の値を変える。

一般化する

Int型だけでなく他の型でも使いたいので、ジェネリクスを利用して一般化した。

func permutation<T>(_ args: [T]) -> [[T]] {
  guard args.count > 1 else { return [args] }
  func rotate(_ args: [T]) -> [T] { [args.last!] + args.dropLast() }
  var rotatedValue = args
  var result: [[T]] = []
  for _ in 0..<args.count {
    let head = rotatedValue.first!
    let tail = rotatedValue.dropFirst().map { $0 as T }
    permutation(tail).forEach {
      result.append([head] + $0)
    }
    rotatedValue = rotate(rotatedValue)
  }
  return result
}

最後に

順列を全列挙する関数を実装した。もっと高速に求める方法があるかもしれないが、ひとまず関数化したことで時短することには成功した。