75. 颜色分类 @Python @Go

原题

https://leetcode.cn/problems/sort-colors/description/

思路

双指针, 分别指向头部和尾部

复杂度

时间：O(n)
空间：O(n)

Python代码

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        p1 = 0 
        p2 = len(nums) - 1 
        temp = nums.copy()
        for i in range(len(temp)):
            if temp[i] == 0:
                nums[p1] = 0 
                p1 += 1
            elif temp[i] == 2:
                nums[p2] = 2
                p2 -= 1 
        for i in range(p1, p2 + 1):
            nums[i] = 1

Go代码

func sortColors(nums []int) {
	// 双指针
	p1, p2 := 0, len(nums)-1
	// 将nums复制到临时切片
	temp := make([]int, len(nums))
	copy(temp, nums)
	for i := 0; i < len(temp); i++ {
		if temp[i] == 0 {
			nums[p1] = 0
			p1++
		} else if temp[i] == 2 {
			nums[p2] = 2
			p2--
		}
	}
	// 中间部分为白色
	for i := p1; i <= p2; i++ {
		nums[i] = 1
	}
}