Interview Prep

Sort Colors Walkthrough

 · 10 mins read

This is a detailed walkthrough of the Sort Colors problem from Leetcode. (Link here). This is based off an email that I sent to my section earlier about how I like to approach technical questions.

Question

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 supposed 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]

Walkthrough of Optimized Solution

I’ll assume that you all understand the problem at the step, so I’m going to talk about my thought process starting from identifying sub-problems / sub-tasks.

I’m just going to be doing a walkthrough of the optimized version, which completes this problem in a single pass and uses constant space. In other words, the time complexity is and the space complexity is , where is the size of the input array.

Identifying Sub-Problems and Sub-Tasks

Since I’m trying to do this in place, I think one sub-task that I will need to perform is a swap in an array. I’ll probably want some helper method to do this since swap code is a little verbose, and I don’t want it to clutter up my main solution:

/**
 * Swaps items in an array.
 * @param a Index of the first element to swap.
 * @param b Index f the second element to swap.
 */
private void swap(int a, int b, int[] nums);

Skeleton & Invariants (the hard part)

My first thought is that I’ll need to iterate through the array somehow to partition all the elements. I’m not exactly sure what my conditions are going to be yet, but I think that somehow I’ll need to keep track of an index for my current location, and that if I’ll need to do something different when I encounter a 0, 1 or 2.

I also want to keep track of my partitions. The first partition is going to be the index of the last 0, and the second partition is going to be the index of the first 2. It doesn’t really matter what my rules for my partitions are, as long as I’m consistent with them.

The idea is that as I’m progressing through nums, I’m going to move the first partition forwards as I encounter 0s, and I’m going to move the second partition backwards as I encounter 2s.

I can’t be sure if there are going to be any 0s in the list, so I’m going to set first to -1. I also can’t be sure if there are any 2s in my list, so I’m going to set second to nums.length.

I’m not exactly sure how I’m handling these cases yet; I just know that I’ll have to handle them eventually, so I’ll write down the code now:

class Solution {
    public void sortColors(int[] nums) {
        int first = -1;             // The index of my last 0.
        int second = nums.length;   // The index of my first 2.
        int index = 0;              // The index that I'm currently exploring.
        while (/* some condition */) {
            if (nums[index] == 0) {
                // Do stuff
            } else if (nums[index] == 1) {
                // Do some other stuff
            } else {
                // Do other stuff
            }
        }
    }

    /**
     * Swaps items in an array.
     * @param a Index of the first element to swap.
     * @param b Index f the second element to swap.
     */
    private void swap(int a, int b, int[] nums) {
        // something here
    }
}

To recap, here are the rules (invariants) I’m making for my solution:

  • first is the index of the last 0 that I have encountered. In other words, all numbers in nums[0..first] are 0.
  • second is the index of the first 2 that I have encountered. In other words, all numbers in nums[second..nums.length-1] are 2.
  • index is the index that we’re currently analyzing. When we entire the while loop, we know that the elements from nums[0..index-1] have been placed correctly.
  • All numbers from nums[first..index-1] are 1. This is probably the trickiest invariant to come up with (but is an important one!). If we know that all elements placed from nums[0..index-1] are correct and that all elements from nums[0..first] are 0, then the remaining elements from nums[first+1..index-1] must be 1.

Filling in the code

I like to fill in my code starting with the helper functions, and work my way from easy problems to hard problems. When I’m writing loops, I also like to write the loop body first (based off the invariants I created in the earlier step), and then write the condition later.

Writing Swap

First: I’m going to write swap. Once I’ve written this helper method, I can forget how it works and just assume that it does what it promises.

    /**
     * Swaps items in an array.
     * @param a Index of the first element to swap.
     * @param b Index f the second element to swap.
     */
    private void swap(int a, int b, int[] nums) {
        int temp = nums[a];
        nums[a] = b;
        nums[b] = temp;
    }

Writing the while loop body

Next, I’m going to tackle each of the cases for when nums[index] = 0, 1, or 2.

  • When nums[index] == 0: I want to swap the 0 that I encountered here to be with the rest of the 0s. I know that all the elements from nums[0..first] are guaranteed to be 0. I don’t know what nums[first + 1] is though (but I know that it’s next to the 0s), so I should swap nums[index] with the element at nums[first + 1]. To maintain the first invariant, I should increase my partition for first because I found a new 0. I should also explore the next index, so I should increment index.
  • When nums[index] == 2: I want to swap the 2 that I encountered here to be with the rest of the 2s. I know that all the elements from nums[second..nums.length-1] are guaranteed to be 2. I don’t know what nums[second - 1] is though (but I know that it’s next to the 2s), so I should swap nums[index] with the element at nums[second - 1]. To maintain the second invariant, I should decrease my partition for second because I found a new 2.
  • When nums[index] == 1: This doesn’t affect any of my first or second invariants, so I’m not going to touch those. However, it does increase our “known” territory, so we should increase index.

This covers all the cases in my while loop.

    while (/* some condition */) {
        if (nums[index] == 0) {
            swap(index, first + 1, nums);
	        first++;
            index++;
        } else if (nums[index] == 1) {
            index++;
        } else {
            swap(index, second - 1, nums);
            second--;
        }
    }

Writing while loop condition

From my invariants earlier, I know that these statements must be true:

  • All elements from nums[0..index] must be placed correctly.
  • All elements from nums[second..nums.length-1] must be placed correctly.

index is the element that we’re using to “explore”. Once index == second, we’re done! When index == second, the following statements are true:

  • All elements from nums[0..second] must be placed correctly.
  • All elements from nums[second..nums.length-1] must be placed correctly.

Therefore, all elements from nums[0..nums.length-1] must be placed correctly, which is exactly what we wanted :)

    while (index < second) {
        ...
    }

Final solution

class Solution {
    public void sortColors(int[] nums) {
        int first = -1;
        int second = nums.length;
        int index = 0;
        
        while (index < second) {
            if (nums[index] == 0) {
                swap(index, first + 1, nums);
                first++;
                index++;
            } else if (nums[index] == 1) {
                index++;
            } else {
                swap(index, second - 1, nums);
                second--;
            }
        }
    }
    
    private void swap(int a, int b, int[] nums) {
        int temp = nums[a];
        nums[a] = nums[b];
        nums[b] = temp;
    }
}