Description

Sort Colors

Given an array with n objects colored red, white or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white and blue.

Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively.

Note: You are not suppose to use the library’s sort function for this problem.

Example:

Input: [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]
Follow up:

  • A rather straight forward solution is a two-pass algorithm using counting sort. First, iterate the array counting number of 0’s, 1’s, and 2’s, then overwrite array with total number of 0’s, then 1’s and followed by 2’s.

  • Could you come up with a one-pass algorithm using only constant space?

对一个只包含0、1、2三个数字的数组进行排序，要求时间复杂度O(n)，空间复杂度O(1)。

Solution

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

func sortColors(nums []int) {
    length := len(nums)
    if length <= 1 {
        return
    }

    l, r := 0, length-1

    // 记录 0（红色）的数量
    red := 0

    for l <= r {
        fmt.Println(nums[l])
        switch nums[l] {
        case 0:
            if l != red {
                // 将第 red 个和第 l 个交换
                // 因为 red 表示 0 的数量，所以第 red 个应该是被跳过的 1
                nums[l], nums[red] = nums[red], nums[l]
            }
            red++
            fallthrough
        case 1:
            // 遇到 1 （白色）则跳过
            l++
        case 2:
            // 遇到 2 （蓝色）则交换到末尾，并且 r--，l 不动
            nums[r], nums[l] = nums[l], nums[r]
            r--
        }
    }
}