75. Sort Colors #

Problem #

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 1:

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?

Problem Statement #

The abstract meaning of the problem is actually sorting. This problem can be solved with quicksort in one pass.

Solution Approach #

The Follow up at the end of the problem proposes a higher requirement: can the problem be solved with a single loop? Since only the three numbers 0, 1, and 2 will appear in this problem, it is also possible to control the order by moving cursors. The specific approach: 0 is placed at the very front, so whenever a 0 is added, 1 and 2 need to be placed. 1 comes before 2, so when adding a 1, 2 also needs to be placed. As for the final 2, only the cursor needs to be moved.

This problem can be solved with counting sort, which is suitable for problems where there are very few numbers to be sorted. Use an array of size 3 to count separately, recording the number of occurrences of 0, 1, and 2. Then arrange 0, 1, and 2 according to their counts. The time complexity is O(n), and the space complexity is O(K). In this problem, K = 3.

This problem can also be solved with a one-pass three-way quicksort. The array is divided into 3 parts: the first part contains all 0s, the middle part contains all 1s, and the last part contains all 2s.

Code #

package leetcode

func sortColors(nums []int) {
	zero, one := 0, 0
	for i, n := range nums {
		nums[i] = 2
		if n <= 1 {
			nums[one] = 1
			one++
		}
		if n == 0 {
			nums[zero] = 0
			zero++
		}
	}
}