Python切片操作行为不一致：旋转数组测试异常排查

我在实现LeetCode旋转数组问题时，碰到了Python切片操作的问题：

• 使用nums[:k+1]时，第二个测试用例通过，但第一个失败，失败后数组变为[99, -1, -100, 3]
• 使用nums[0:k]时，第一个测试用例通过，但第二个失败，失败后数组变为[4, 5, 6, 7, 1, 2, 3]

相关代码 #

rotate.py

def rotate(nums, k):
    temp = nums[0:k+1]
    del nums[0:k+1]
    nums.extend(temp)

test_rotate.py

import unittest
from rotate_array import rotate

class TestRotateArray(unittest.TestCase):
    def test_rotate_array2(self):
        nums= [-1, -100, 3, 99]
        k = 2
        expected = [3, 99, -1, -100]
        rotate(nums, k)
        self.assertListEqual(nums, expected)
    
    def test_rotate_array(self):
        nums= [1,2,3,4,5,6,7]
        k = 3
        expected = [5,6,7,1,2,3,4]
        rotate(nums, k)
        self.assertListEqual(nums, expected)

if __name__ == '__main__':
        unittest.main()

问题原因 #

核心错误是搞反了旋转操作的元素范围：题目要求向右旋转k步，本质是把数组的最后k个元素移动到数组开头，而不是操作数组开头的元素。

比如第一个测试用例，数组[1,2,3,4,5,6,7]向右旋转3步，需要把最后3个元素[5,6,7]移到前面；而你之前操作前k或k+1个元素，相当于做了向左旋转，结果自然不符合预期。

修正后的代码 #

正确的实现需要先处理k大于数组长度的情况（比如数组长度为4，k=5时，等价于旋转1步），然后移动末尾k个元素到开头：

def rotate(nums, k):
    n = len(nums)
    k = k % n  # 处理k超过数组长度的情况
    # 把最后k个元素移到开头
    temp = nums[-k:]
    del nums[-k:]
    nums[:0] = temp

或者更简洁的原地修改写法：

def rotate(nums, k):
    n = len(nums)
    k %= n
    nums[:] = nums[-k:] + nums[:-k]

这个版本能同时通过两个测试用例，符合题目要求的原地修改数组的规则。

内容的提问来源于stack exchange，提问作者Jamal