leetcode75：颜色分类

题目链接

leetcode

题目描述

给定一个包含红色、白色和蓝色、共n个元素的数组nums ，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。
整数0、 1 和 2分别表示红色、白色和蓝色。
在不使用库内置的 sort 函数的情况下解决这个问题。

C++ 代码

#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;

void printArray(const vector<int>& nums) {
    for (int num : nums) {
        cout << num << " ";
    }
    cout << endl;
}

/*
    快速排序的变形
    针对这道题有更优的解法，但是快排可以扩展到任意多种颜色。

    时间复杂度：O(n)
        其中 n 是数组的长度，k 是数组中不同颜色的数量。
        对于这个问题，时间复杂度是 O(nlogk)。
        k = 3，由于 log3 是一个常数，所以时间复杂度可以简化为 O(n)。
    空间复杂度：O(logk)
        递归的深度是 logk，因此空间复杂度是 O(logk)。
        由于 k = 3，空间复杂度是 O(log3)，这可以简化为 O(1)。
*/
class Solution_0 {
public:
    void sortColors(vector<int>& nums) {
        partition(nums, 0, nums.size() - 1, 0, 2);
    }

    void partition(vector<int>& nums, int start, int end, int lower, int upper) {
        if (start >= end || lower >= upper) return;

        int pivot = lower + (upper - lower) / 2;
        int l = start, r = end;
        while (l <= r) {
            while (l <= r && nums[l] <= pivot) l ++;
            while (l <= r && nums[r] > pivot) r --;
            if (l <= r) swap(nums[l ++], nums[r --]);
        }

        partition(nums, start, r, lower, pivot);
        partition(nums, l, end, pivot + 1, upper);
    }
};

/*
    参考：三路快排，荷兰国旗
    题目规定只有 3 种颜色，只需要完成一次类似快排 partition 的操作。

    时间复杂度：O(n)
    空间复杂度：O(1)
*/
class Solution_1 {
public:
    void sortColors(vector<int>& nums) {
        for (int i = 0, j = 0, k = nums.size() - 1; i <= k; ) {
            if (nums[i] == 0) swap(nums[i ++], nums[j ++]);
            else if (nums[i] == 1) i ++;
            else if (nums[i] == 2) swap(nums[i], nums[k --]);
        }
    }
};

/*
    计数排序
    统计 3 种颜色各自的个数，这个做法可以推广到多个颜色。

    时间复杂度：O(n)
        其中 n 是数组的长度，k 是数组中不同颜色的数量。
        遍历了两次数组，时间复杂度是 O(n)。
    空间复杂度：O(1)
        需要 k 个元素记录颜色的数量，所以是 O(k)。
        由于 k = 3，空间复杂度是 O(log3)，这可以简化为 O(1)。
*/
class Solution {
public:
    void sortColors(vector<int>& nums) {
        int n = nums.size();
        int cnt0 = 0, cnt1 = 0, cnt2 = 0;
        for (const int& num : nums) {
            if (num == 0) cnt0 ++;
            else if (num == 1) cnt1 ++;
            else if (num == 2) cnt2 ++;
        }

        for (int i = 0; i < n; i ++) {
            if (i < cnt0) nums[i] = 0;
            else if (i - cnt0 < cnt1) nums[i] = 1;
            else nums[i] = 2;
        }
    }
};

int main() {
    Solution solution;
    vector<vector<int>> nums_cases = {
        {2,0,2,1,1,0},
        {2,0,1}
    };

    for (int i = 0; i < nums_cases.size(); i ++) {
        vector<int> nums = nums_cases[i];
        cout << "Input: ";
        printArray(nums);
        solution.sortColors(nums);
        cout << "Output: ";
        printArray(nums);
    }
    
    return 0;
}

Golang 代码

package main

import (
	"fmt"
)

/*
   三指针方法：
       使用三个指针：left、i 和 right。
       left 指针用于记录 0 的边界。
       i 指针用于遍历数组。
       right 指针用于记录 2 的边界。
   状态转移：
       如果当前元素是 0，与 left 指针交换，并将 left 和 i 同时向右移动。
       如果当前元素是 1，直接将 i 向右移动。
       如果当前元素是 2，与 right 指针交换，并将 right 向左移动。
   边界条件：
       如果数组长度小于等于 1，直接返回。

   时间复杂度：O(n)
       遍历一次数组，时间复杂度为 O(n)，其中 n 是数组的长度。
   空间复杂度：O(1)
       使用了常数级别的额外空间，空间复杂度为 O(1)。
*/
// sortColors 对颜色数组进行排序
func sortColors(nums []int) {
	n := len(nums)
	if n <= 1 {
		return
	}

	// 初始化指针
	left, right := 0, n - 1
	i := 0

	// 使用三指针方法进行排序
	for i <= right {
		switch nums[i] {
		case 0:
			// 如果当前元素是 0，与 left 指针交换
			nums[i], nums[left] = nums[left], nums[i]
			left ++
			i ++
		case 1:
			// 如果当前元素是 1，跳过
			i ++
		case 2:
			// 如果当前元素是 2，与 right 指针交换
			nums[i], nums[right] = nums[right], nums[i]
			right --
		}
	}
}

func main() {
	// 测试用例
	testCases := []struct {
		nums    []int
		expected []int
	}{
		{
			nums:    []int{2, 0, 2, 1, 1, 0},
			expected: []int{0, 0, 1, 1, 2, 2},
		},
		{
			nums:    []int{2, 0, 1},
			expected: []int{0, 1, 2},
		},
	}

	for i, tc := range testCases {
		sortColors(tc.nums)
		fmt.Printf("Test Case %d, Input: nums = %v\n", i+1, tc.nums)

		if fmt.Sprint(tc.nums) == fmt.Sprint(tc.expected) {
			fmt.Printf("Test Case %d, Output: %v, PASS\n", i+1, tc.nums)
		} else {
			fmt.Printf("Test Case %d, Output: %v, FAIL (Expected: %v)\n", i+1, tc.nums, tc.expected)
		}
	}
}