DEV Community

loading...

LeetCode - Sort Colors

Alkesh Ghorpade
Ruby on Rails and Golang consultant
Originally published at alkeshghorpade.me ・4 min read

Container

Problem statement

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.

Problem statement taken from: https://leetcode.com/problems/sort-colors

Example 1:

Input: nums = [2, 0, 2, 1, 1, 0]
Output: [0, 0, 1, 1, 2, 2]
Enter fullscreen mode Exit fullscreen mode

Example 2:

Input: nums = [2, 0, 1]
Output: [0, 1, 2]
Enter fullscreen mode Exit fullscreen mode

Example 3:

Input: nums = [0]
Output: [0]
Enter fullscreen mode Exit fullscreen mode

Example 4:

Input: nums = [1]
Output: [1]
Enter fullscreen mode Exit fullscreen mode

Constraints:

- n == nums.length
- 1 <= n <= 300
- nums[i] is 0, 1, or 2
Enter fullscreen mode Exit fullscreen mode

Explanation

Simple counting

The simple approach will be to count the occurrences of each integer 0, 1
and 2 using three different variables.

Using the above three count variables we fill in the array.

A small C++ snippet of the approach will look like this:

for (int i = 0; i < n; i++) {
    if (arr[i] == 0)
        count0++;
    if (arr[i] == 1)
        count1++;
    if (arr[i] == 2)
        count2++;
}

for (int i = 0; i < count0; i++)
    arr[i] = 0;

for (int i = count0; i < (count0 + count1); i++)
    arr[i] = 1;

for (int i = (count0 + count1); i < n; i++)
    arr[i] = 2;
Enter fullscreen mode Exit fullscreen mode

The time complexity of the above program is O(N).
But in the above approach, we iterate the array twice.

Dutch national flag problem

We can use the approach of
Dutch national flag problem.
The problem was posed with three colors,
here we have 0, 1, and 2.
The array is divided into four sections to solve this problem.

Let's check the algorithm:

-  Keep three indices low = 1, mid = 1 and, high = N and, there are four ranges,
   1 to low (the range containing 0),
   low to mid (the range containing 1),
   mid to high (the range containing unknown elements) and
   high to N (the range containing 2).

-  Traverse the array from start to end and mid is less than high. (Loop counter is i)

-  if element == 0
  - swap the element with the element at index low and update low = low + 1 and mid = mid + 1

-  if element == 1
  - set mid = mid + 1

-  if element == 2
  - swap the element with the element at index high and update high = high – 1.
  - set i = i – 1.

- return array
Enter fullscreen mode Exit fullscreen mode

The time complexity of the program is O(N) as we iterate the array only once.
Space complexity is O(1) because we do not use any other additional data structures.

C++ solution
class Solution {
public:
    void sortColors(vector<int>& nums) {
        int low = 0, mid = 0, high = nums.size() - 1;

        while(mid <= high){
            switch (nums[mid]){
                case 0:
                    swap(nums[low++], nums[mid++]);
                    break;
                case 1:
                    mid++;
                    break;
                case 2:
                    swap(nums[mid], nums[high--]);
                    break;
            }
        }
    }
};
Enter fullscreen mode Exit fullscreen mode
Golang solution
func sortColors(nums []int)  {
    low := 0
    mid := 0
    high := len(nums) - 1

    for mid <= high {
        switch (nums[mid]) {
            case 0:
                tmp := nums[low]
                nums[low] = nums[mid]
                nums[mid] = tmp
                low++
                mid++
                break
            case 1:
                mid++
                break
            case 2:
                tmp := nums[mid]
                nums[mid] = nums[high]
                nums[high] = tmp
                high--
                break
        }
    }
}
Enter fullscreen mode Exit fullscreen mode
Javascript solution
var sortColors = function(nums) {
    function swap(i, j) {
        [nums[i], nums[j]] = [nums[j], nums[i]];
    }

    let low = 0;
    let high = nums.length - 1;
    let mid = 0;

    while (mid <= high) {
    const n = nums[mid];

      if (n === 0) {
        swap(mid, low);
        low++;
        mid++;
      } else if (n === 2) {
        swap(mid, high);
        high--;
      } else {
        mid++;
      }
    }
};
Enter fullscreen mode Exit fullscreen mode

Let's dry run the problem

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

Step 1: low = 0
        mid = 0
        high = nums.length() - 1
             = 6 - 1
             = 5

Step 2: loop while mid < = high
        0 <= 5
        true

        switch (nums[mid])
        nums[mid] = nums[0]
                  = 2

        case 2:
        swap(nums[mid], nums[high--])
        swap(nums[0], nums[5])
        swap(2, 0)

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

        high--
        high = 4

Step 3: loop while mid < = high
        0 <= 4
        true

        switch (nums[mid])
        nums[mid] = nums[0]
                  = 0

        case 0:
        swap(nums[low++], nums[mid++])
        swap(nums[0], nums[0])
        swap(0, 0)

        nums = [0, 0, 2, 1, 1, 2]
        low++
        mid++

        low = 1
        mid = 1

Step 4: loop while mid < = high
        1 <= 4
        true

        switch (nums[mid])
        nums[mid] = nums[1]
                  = 0

        case 0:
        swap(nums[low++], nums[mid++])
        swap(nums[1], nums[1])
        swap(1, 1)

        nums = [0, 0, 2, 1, 1, 2]
        low++
        mid++

        low = 2
        mid = 2

Step 5: loop while mid < = high
        2 <= 4
        true

        switch (nums[mid])
        nums[mid] = nums[2]
                  = 2

        case 2:
        swap(nums[mid], nums[high--])
        swap(nums[2], nums[4])
        swap(2, 1)

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

        high--
        high = 3

Step 6: loop while mid < = high
        2 <= 3
        true

        switch (nums[mid])
        nums[mid] = nums[2]
                  = 1

        case 1:
        mid++
        mid = 3

Step 7: loop while mid < = high
        3 <= 3
        true

        switch (nums[mid])
        nums[mid] = nums[3]
                  = 1

        case 1:
        mid++
        mid = 4

Step 8: loop while mid < = high
        4 <= 3
        false

The result is [0, 0, 1, 1, 2, 2]
Enter fullscreen mode Exit fullscreen mode

Discussion (0)