Skip to content

75. Sort Colors

Description

You are given an array nums 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.

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

You must solve this problem without using the library's sort function.

 

Example 1:

Input: nums = [2,0,2,1,1,0]

Output: [0,0,1,1,2,2]

Explanation:

The array has two 0s, two 1s, and two 2s. Sorting them in-place places all 0s first, then all 1s, then all 2s.

Example 2:

Input: nums = [2,0,1]

Output: [0,1,2]

Explanation:

The array has one each of 0, 1, and 2, arranged in-place in the order 0, 1, 2.

 

Constraints:

  • n == nums.length
  • 1 <= n <= 300
  • nums[i] is either 0, 1, or 2.

 

Follow up: Could you come up with a one-pass algorithm using only constant extra space?

Solutions

Solution 1: Three Pointers

Thinking

Count \(0,1,2\) then overwrite: two passes. Sorting would also work for \(n \le 300\). The problem forbids library sort, and the follow-up wants one scan and \(O(1)\) space.

Three values only need a partition into all \(0\)s, all \(1\)s, and all \(2\)s. Let \(i\) and \(j\) bound the \(0\)s and \(2\)s already placed, and let \(k\) scan the unknown middle. A swap toward \(2\) brings an unseen value, so \(k\) stays; a swap toward \(0\) brings a value from the scanned range, so \(k\) advances too. One pass finishes the three segments.

We define three pointers \(i\), \(j\), and \(k\). Pointer \(i\) is used to point to the rightmost boundary of the elements with a value of \(0\) in the array, and pointer \(j\) is used to point to the leftmost boundary of the elements with a value of \(2\) in the array. Initially, \(i=-1\), \(j=n\). Pointer \(k\) is used to point to the current element being traversed, initially \(k=0\).

When \(k < j\), we perform the following operations:

  • If \(nums[k] = 0\), then swap it with \(nums[i+1]\), then increment both \(i\) and \(k\) by \(1\);
  • If \(nums[k] = 2\), then swap it with \(nums[j-1]\), then decrement \(j\) by \(1\);
  • If \(nums[k] = 1\), then increment \(k\) by \(1\).

After the traversal, the elements in the array are divided into three parts: \([0,i]\), \([i+1,j-1]\) and \([j,n-1]\).

The time complexity is \(O(n)\), where \(n\) is the length of the array. Only one traversal of the array is needed. The space complexity is \(O(1)\).

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution:
    def sortColors(self, nums: List[int]) -> None:
        i, j, k = -1, len(nums), 0
        while k < j:
            if nums[k] == 0:
                i += 1
                nums[i], nums[k] = nums[k], nums[i]
                k += 1
            elif nums[k] == 2:
                j -= 1
                nums[j], nums[k] = nums[k], nums[j]
            else:
                k += 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
    public void sortColors(int[] nums) {
        int i = -1, j = nums.length, k = 0;
        while (k < j) {
            if (nums[k] == 0) {
                swap(nums, ++i, k++);
            } else if (nums[k] == 2) {
                swap(nums, --j, k);
            } else {
                ++k;
            }
        }
    }

    private void swap(int[] nums, int i, int j) {
        int t = nums[i];
        nums[i] = nums[j];
        nums[j] = t;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
class Solution {
public:
    void sortColors(vector<int>& nums) {
        int i = -1, j = nums.size(), k = 0;
        while (k < j) {
            if (nums[k] == 0) {
                swap(nums[++i], nums[k++]);
            } else if (nums[k] == 2) {
                swap(nums[--j], nums[k]);
            } else {
                ++k;
            }
        }
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
func sortColors(nums []int) {
    i, j, k := -1, len(nums), 0
    for k < j {
        if nums[k] == 0 {
            i++
            nums[i], nums[k] = nums[k], nums[i]
            k++
        } else if nums[k] == 2 {
            j--
            nums[j], nums[k] = nums[k], nums[j]
        } else {
            k++
        }
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
/**
 Do not return anything, modify nums in-place instead.
 */
function sortColors(nums: number[]): void {
    let i = -1;
    let j = nums.length;
    let k = 0;
    while (k < j) {
        if (nums[k] === 0) {
            ++i;
            [nums[i], nums[k]] = [nums[k], nums[i]];
            ++k;
        } else if (nums[k] === 2) {
            --j;
            [nums[j], nums[k]] = [nums[k], nums[j]];
        } else {
            ++k;
        }
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
impl Solution {
    pub fn sort_colors(nums: &mut Vec<i32>) {
        let mut i = -1;
        let mut j = nums.len();
        let mut k = 0;
        while k < j {
            if nums[k] == 0 {
                i += 1;
                nums.swap(i as usize, k as usize);
                k += 1;
            } else if nums[k] == 2 {
                j -= 1;
                nums.swap(j, k);
            } else {
                k += 1;
            }
        }
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
public class Solution {
    public void SortColors(int[] nums) {
        int i = -1, j = nums.Length, k = 0;
        while (k < j) {
            if (nums[k] == 0) {
                swap(nums, ++i, k++);
            } else if (nums[k] == 2) {
                swap(nums, --j, k);
            } else {
                ++k;
            }
        }
    }

    private void swap(int[] nums, int i, int j) {
        int t = nums[i];
        nums[i] = nums[j];
        nums[j] = t;
    }
}

Comments