【leetcode刷题】颜色分类——双指针遍历排序

前言

题目：给定一个包含红色、白色和蓝色、共 n 个元素的数组 nums ，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。要求不能使用sort函数。
思路：我们设置双指针p1=0，p2=len(nums)-1，从开头i=0开始遍历数组用while循环判断nums[i]是否为0,1,2这三个数，0则交换nums[i]和nums[p1]，且i和p1都向前移动一步，1则i向前移动，2则交换p2以此类推。不用for循环是因为我们交换完数字后还需要继续判断当前nums[i]的值进行相应的操作。

代码如下（python）：

class Solution:
    def sortColors(self, nums):
        def swap(nums, a, b):
            nums[a], nums[b] = nums[b], nums[a]
        p1 = 0
        p2 = len(nums)-1
        i = 0
        while i < p2:
            if nums[i] == 0:
                swap(nums, i, p1)
                i += 1
                p1 += 1
            elif nums[i] == 1:
                i += 1
            else:
                swap(nums, i, p2)
                p2 -= 1
        return nums

a = Solution()
nums = list(map(int,input().split()))
print(a.sortColors(nums))t

python交换元素可直接x,y=y,x不用加中间变量temp

示例