オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

LeetCode / Rotate Array

https://leetcode.com/problems/rotate-array/

Given an array, rotate the array to the right by k steps, where k is non-negative.

Example 1:
Input: [1,2,3,4,5,6,7] and k = 3
Output: [5,6,7,1,2,3,4]

Explanation:
rotate 1 steps to the right: [7,1,2,3,4,5,6]
rotate 2 steps to the right: [6,7,1,2,3,4,5]
rotate 3 steps to the right: [5,6,7,1,2,3,4]

Example 2:
Input: [-1,-100,3,99] and k = 2
Output: [3,99,-1,-100]

Explanation:
rotate 1 steps to the right: [99,-1,-100,3]
rotate 2 steps to the right: [3,99,-1,-100]

Note:
Try to come up as many solutions as you can, there are at least 3 different ways to solve this problem.
Could you do it in-place with O(1) extra space?

Rotate Arrayということで直訳すると配列を回転させろ、ということですが、インデックスを前後にずらす操作を行う問題です。
解法を最低3つ、さらに空間計算量はO(1)で、と言われています。AND条件ではなくOR条件だと思いますが、意外に難しく感じました。

解答・解説

解法1

私の力ではどうしても空間計算量がO(n)になってしまいました。

インデックスを前後にずらす=元のリストを2つのリストに分割し組み替えて再結合することになるので、以下コードで正しい値が得られます。

class Solution:
    def rotate(self, nums: List[int], k: int) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        k %= len(nums)
        nums[:] = nums[-k:]+nums[:-k]

今回はnumsそのものを操作して戻り値にするという指示があるため、リストに対して破壊的な操作をするためにnums[:]としてコピーを作成します。

リスト内包表記で書くこともできます。

class Solution:
    def rotate(self, nums: List[int], k: int) -> None:
        k %= len(nums)
        nums[:] = [nums[i] for i in range(-k,len(nums)-k)]
解法2

空間計算量がO(1)の解法その1。
以下Solutionの転載ですが、

Original List : 1 2 3 4 5 6 7
After reversing all numbers : 7 6 5 4 3 2 1
After reversing first k numbers : 5 6 7 4 3 2 1
After revering last n-k numbers : 5 6 7 1 2 3 4 --> Result

3つの手順から成り、1.リスト全体を逆順にソートし、2.リスト後方のk個の要素だけで逆順にソートし、3.残りのn-k個の要素だけで逆順にソートするという手法です。

class Solution:
    def rotate(self, nums: List[int], k: int) -> None:
        k %= len(nums)
        self.reverse(nums, 0, len(nums)-1)
        self.reverse(nums, 0, k-1)
        self.reverse(nums, k, len(nums)-1)
    def reverse(self, nums, start, end):
        while start < end:
            tmp = nums[start]
            nums[start] = nums[end]
            nums[end] = tmp
            start += 1
            end -= 1

空間計算量を抑えた解法の中で、これが最も明快だと思います。

解法3

空間計算量がO(1)の解法その2。しかしこれは難しいと思いました。

手順は大きく分けると2つから成り、nums後方のk個の要素を正しい位置に入れ替え、次に残りのnums前方のn-k個の要素を正しい位置に入れ替えます。
入れ替え対象のリストは初めはnums全体ですが、入れ替えが完了するまで徐々に狭まっていきます。

具体例で示します。以下、
n: 入れ替え対象のリストの要素数
k: インデックスをずらす要素数
j: 入れ替え対象リストの始点インデックス
とします。つまり常に n + j == len(nums) となります。

nums = [1, 2, 3, 4, 5, 6, 7] を例に考えます。このとき、n = 7, k = 3, j = 0 です。
初めの入れ替えで、以下のように進みます。

[5, 2, 3, 4, 1, 6, 7]
[5, 6, 3, 4, 1, 2, 7]
[5, 6, 7, 4, 1, 2, 3]

ここまで来て、さらに入れ替える必要があるのは後方の [4, 1, 2, 3] です。
このとき、n = n - k, j = j + k, k %= n と計算し、n = 4, k = 3, j = 3 です。

次の入れ替えは以下のように進みます。

[5, 6, 7, 1, 4, 2, 3]
[5, 6, 7, 1, 2, 4, 3]
[5, 6, 7, 1, 2, 3, 4]

これで、入れ替え完了です。
先ほどと同様に、n = n - k, j = j + k, k %= nと計算し、n = 1, k = 0, j = 6 です。

入れ替え完了の判定方法は、n <= 0 と k % n == 0 の2パターンがあります。
n <= 0 になると、入れ替え対象のリストが存在しないので終了し、
k % n == 0 になると、インデックスをずらしても1周回って同じリストとなってしまうので、終了します。

以上をコードに落とすと、以下のようになります。

class Solution(object):
    def rotate(self, nums, k):
        n, k, j = len(nums), k % len(nums), 0
        while n > 0 and k % n != 0:
            for i in range(0, k):
                nums[j + i], nums[len(nums) - k + i] = nums[len(nums) - k + i], nums[j + i] # swap
            n, j = n - k, j + k
            k = k % n

初見でこの解法にたどり着く人は神なんじゃないかな。