Thinking About Partitions

In algorithmic problems, the concept of array partitions comes up too often for us to ignore.

Thus far, we’ve looked at some pretty complex challenges such as Minimum Palindromic Partitions and Partitioning an Array into K Equal Sum Subsets. We used Dynamic Programming to attain answers for both partitioning problems.

In this article, I’d like to explore more “bread and butter” use cases for dividing arrays into subarrays. We’ll kick off our discussion with our problem of the day, the Color Sorting problem (also known as the Dutch Partition problem).

The Color Sorting Problem

From LeetCode:

Given an array with n objects colored red, white, or blue, sort them in-place in a single pass. Objects of the same color should be adjacent to each other and the sorted colors should be in the order of red, white, and blue.

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

From the outset, we know that we’ll need to store some kind of state in the array while doing our linear traversal –– we don’t have extra space to work with since the question expects us to shift elements in-place.

A reasonable approach would be to keep track of three color buckets (for red, white, and blue). We can represent these three buckets as three subarrays in our input array. In addition, we implicitly maintain a fourth subarray which contains unbucketed / uncategorised elements.

We use index variables to mark the boundaries of our color buckets. A mid-traversal snapshot of these indices might look something like this:

# A[:r] -> Red bucket
# A[r:w] -> White bucket
# A[b:] -> Blue bucket
# A[w:b] -> Unseen bucket

"""
Input:   [0, 0, 1, 1, 0, 2, 1, 1, 2, 2]
Pointers:       r     w           b
"""

Notice that I’ve flushed the blue bucket to the right side of the array (i.e. A[b:]). While it’s feasible to lump all three color buckets together on the left side of the array (in the form of A[:r] + A[r:w] + A[w:b]), handling newly seen elements becomes trickier because additions to the red bucket would require us to shift both the white and blue buckets.

Let’s now see how we might grow these buckets as we pass through the array:

def sort_colors(A):
    # A[:red] -> Red bucket
    # A[red:white] -> White bucket
    # A[blue:] -> Blue bucket
    # A[white:blue] -> Unseen bucket
    red = 0
    white = 0
    blue = len(A)

    # Navigate through the array until size of 
    # unseen subarray shrinks to zero
    while white < blue:
        if A[white] == 0:
            A[white], A[red] = A[red], A[white]
            red += 1
            white += 1
        
        elif A[white] == 1:
            white += 1
        
        else: # A[white] == 2
            blue -= 1
            A[white], A[blue] = A[blue], A[white]

A key point here is that white is the main index variable we use for moving through the array and inspecting elements. For every A[white], we check its color and swap it with the boundary element of the corresponding subarray.

Given an array of n numbers, our algorithm runs in O(1) space and O(n) time.

With that, we’ve solved the Color Sorting problem 🔴⚪️🔵.

Though we’ve solved our problem of the day, I’d like to take the chance to mention a couple of other applications of array partitioning.

Dealing with Wider Ranges of Partition-by Values

Not all arrays are going to be in red, white, and blue. What if we don’t know the exact range of elements we’re about to encounter?

To adapt a question from EPI: Suppose we’re given an array of tuples (x, y) and we’re told to partition elements by their x value. We don’t know what these x values are, but we’re told that there are k possibilities for x (and k is a relatively small number). Forget about the single pass constraint.

To start, we can go through the input array to count instances of x and define offset indices for them:

def partition_xy(A):

    x_counts = collections.Counter([x for x, y in A])
    x_offsets = {}
    total_offset = 0
    for x, count in x_counts.items():
        x_offsets[x] = total_offset
        total_offset += count

    # ...

Each offset index demarcates a subarray A[offset:offset+count] that will be responsible for storing all (x, y) tuples that have a specific value of x.

We move (x, y) tuples to their correct position by repeatedly:

• Selecting a random index in A that hasn’t been fitted with the correct x value yet
• Moving the element at the randomly-selected index to the subarray it belongs to

def partition_xy(A):

    # Get x_offsets and x_counts

    while x_offsets:
        # Pick a random x value, x1
        x1 = next(iter(x_offsets))
        x1_offset = x_offsets[x1]

        # Find (x2, y2), situated at the
        # current offset for x1
        x2, y2 = A[offset]

        # Place (x2, y2) in its correct subarray,
        # move the element in (x2, y2)'s position
        # to x1's offset
        x2_offset = x_offsets[x2]
        A[x2_offset], A[x1_offset] = A[x1_offset], A[x2_offset]

        # Increment x2's offset
        x_counts[x2] -= 1
        if x_counts[x2] > 0:
            x_offsets[x2] += 1
        else:
            del x_offsets[x2]

    return A

This effectively allows us to partition an array with many repeats in O(k) space and O(n) time.

Quicksort-style Partitioning

The partitioning algorithm that we wrote for the Color Sorting problem is, in some ways, similar to the one we use for traditional quicksort operations:

def qsort(A, start, last):
    if start < last:
        p = partition(A, start, last)
        qsort(A, start, p-1)
        qsort(A, p+1, last)

def partition(A, start, last):
    p = start

    # A[last] is our pivot value; compare
    # all other elements against A[last]
    for i in range(start, last):

        # Elements less than the pivot
        # will be moved to its left
        if A[i] < A[last]:
            A[i], A[p], A[p], A[i]
            p += 1
    A[last], A[p] = A[p], A[last]

    return p

You might think of sort_colors() and partition() as variants of the same fundamental idea. Both functions scan an array from left to right and group values into subarrays within the same pass. partition(), as it is used in quicksort, compares elements to a single pivot value and divides arrays into two such that elements in the left subarray are smaller than the right subarray. sort_colors(), on the other hand, compares elements to a range of values (A[white] == 0, A[white] == 1, etc).

It’s worth pointing out that you can modify partition() and apply it to questions like Kth Largest Element in array.

def find_kth_largest(A, k):
    def ptn(start, last, k):
        p = start

        # Randomly select pivot, place it at last index
        pivot_idx = random.randint(start, last)
        A[pivot_idx], A[last] = A[last], A[pivot_idx]
        
        for i in range(start, last):
            if A[i] > A[last]:
                A[p], A[i] = A[i], A[p]
                p += 1
        
        A[p], A[last] = A[last], A[p]
        
        # Recurse based on the partition index
        if k - 1 == p:
            return A[p]

        elif k - 1 > p:
            return ptn(p+1, last, k)
        else: # k - 1 < p
            return ptn(start, p-1, k)
    return ptn(0, len(A)-1, k)

Unlike in quicksort, the partition function in find_kth_largest() shifts larger elements to the left of the partition, not the right. We continually recurse on left / right subarrays until our partition index p lies exactly at k - 1. When it does, we would have found the kth largest element at A[p].