今日も窓辺でプログラム

外資系企業勤めのエンジニアが勉強した内容をまとめておくブログ

LeetCode 31: Next Permutation

順列苦手で解けなかったのでメモ。問題はLeetCodeに掲載されていたものを使用します。

LeetCode自体の紹介はこちら。
LeetCode: コーディング面接に向けた練習に使えるサイトの紹介 - 今日も窓辺でプログラム

問題

適当な和訳

整数の配列が入力として与えられるので、その整数を並び替えてできる順列のうち、辞書順で次に来るものを求めよ。辞書順で一番最後に来るものが入力として与えられた場合は、辞書順の最初のものを返すこと。
また、新たにメモリは使わず与えられた配列内で並び替えをすること。

https://leetcode.com/problems/next-permutation/

要は、配列が与えられるので、その配列を次の順列に並び替えろ、という問題です。

入力が[1, 2, 3]の場合は、この配列を[1, 3, 2]に並べ替えないといけません。
[4, 1, 3, 7, 6, 5, 2]の場合は、[4, 1, 5, 2, 3, 6, 7]が答えになります。



アルゴリズム

nums = [4, 1, 3, 7, 6, 5, 2]
という配列が入力として与えられた場合を例に考えます。

ステップ1

まずは、nums[k] < nums[k + 1]という条件が満たされるような最大のkを探します。つまり、配列の後ろから要素がどこまで逆順に並んでいるか調べます。
例の場合だと、nums[2] = 3 より後(= [7, 6, 5, 4, 2])は要素が逆順にソートされた状態になっているので、k = 2となります。
このようなnums[k]が、配列を左から見ていったときに繰り上がる最初の値となります。

このようなnums[k]が存在しない場合は、配列全体が逆順にソートされている状態なので、全体を反転させてあげて辞書順の最初のものを返してあげればOKです。

ステップ2

では、nums[k]が繰り上がらなければならないのはわかりましたが、どのような値に繰り上がればいいのでしょうか。順列の次に来る配列を求めたいので、繰り上がる幅は最小にしたいです。つまり、nums[k+1], nums[k+2], ..., nums[n-1]の中で、nums[k]の次に大きな値にしたいです。
そのような値をnums[l]とすると、lは

nums[l] > nums[k]を満たすような最大のl

ということができます。ステップ1のkの条件から、lが存在してl >= k + 1となることはは保証されています。
nums[k]をnums[l]に繰り上げたかったので、この2つの値を入れ替えてあげます。

例の場合だと

  • nums[k] = nums[2] = 3
  • nums[l] = nums[5] = 5

となるので、入れ替えた後の途中経過は
nums = [4, 1, 5, 7, 6, 3, 2]
となります。

ステップ3

最後に、繰り上がった部分より右側を処理します。nums[k+1]からnums[n-1]まではステップ1, 2を通して逆順でソートされたままになっているので、この部分を反転させて、辞書順で一番小さい順番にしてあげます。
例の場合、繰り上がった箇所の右側であるnums[3]からnums[6]、つまり[7, 6, 3, 2]の部分が逆順のままなので、ここを反転させてあげます。すると
nums = [4, 1, 5, 2, 3, 6, 7]
となり、最初の入力の次の順列が得られていることがわかります。

コード

Pythonのコードです。

class Solution(object):
    def nextPermutation(self, nums):
        """
        :type nums: List[int]
        :rtype: void Do not return anything, modify nums in-place instead.
        """
        # ステップ1:kを求める
        k = -1
        i = len(nums) - 2
        while i >= 0:
            if nums[i] < nums[i + 1]:
                k = i
                break
            i -= 1

        # kがなかったら、配列は逆順にソートされているので、反転して終了
        if k == -1:
            nums.reverse()
            return

        # ステップ2:lを求め、nums[k]とnums[l]を入れ替える
        l = k + 1
        i = len(nums) - 1
        while i > k + 1:
            if nums[k] < nums[i]:
                l = i
                break
            i -= 1

        tmp = nums[k]
        nums[k] = nums[l]
        nums[l] = tmp

        # ステップ3:nums[k]より右側の要素を反転させる
        nums[k+1:] = reversed(nums[k+1:])

(コードを簡潔にするために最後はreversed関数を使ってますが、本当はここはスワップを使って余計なメモリを使わずに実装しなければいけません。)