[EPI] 5.1 Dutch National Flag Problem

4 minute read

Dutch National Flag Problem - Elements of Programming Interviews (EPI)

Problem:

Write a program that takes an array \(A\) and an index \(i\) and rearranges \(A\) such that all elements less than the pivot: \(A[i]\) appear first, followed by elements equal to the pivot, followed by elements greater than the pivot. [4,4,2,1,1,1,3,5,0] => if i = 1, pivot = 4 output should be [2,1,1,1,0,3,4,4,5] or equivalent.

Solution:

The naive approach would be to make 3 separate lists (less, equal and greater) and iterate over \(A\) and compare each element to \(A[i]\) and based on the comparison append it to the appropriate list and in the end combine the lists and set \(A\) equal to the result. This approach has \(O(n)\) time complexity but also uses \(O(n)\) extra space.

We can reduce the extra space requirement at the cost of time complexity by making multiple passes over \(A\) and swapping smaller elements to the start of the array in the first pass and swapping larger elements to the end of the array in the second pass.

We can further reduce the runtime by doing everything in one pass. This can be done by partitioning into 4 subarrays:

A[:lt] : elements less than pivot.
A[lt:equal] : elements equal to pivot.
A[equal:gt] : elements not yet classified.
A[gt:] : elements greater than pivot.

Here lt (less than), equal, gt (greater than) are indexes into the array which keep track of the boundaries of the subarrays. We keep these invariants during the algorithm.

def dutch_flag_partition(pivot_index: int, A: List[int]) -> None:
    """Function to reorder an array A into 4 parts, part1: less
    than pivot index element, part 2 equal to pivot index element,
    part 3 unclassified, part4: greater than pivot index"""
    lt, equal, gt = (0, 0, len(A)-1)
    p = A[pivot_index]
    while equal <= gt:
        if A[equal] < p:
            A[lt], A[equal] = A[equal], A[lt]
            lt += 1
            equal += 1
        elif A[equal] > p:
            A[gt], A[equal] = A[equal], A[gt]
            gt -= 1
        elif A[equal] == p:
            equal += 1

Variants:

Let keys take one of 3 values. Reorder the array such that all objects with the same key appear together. The order of the subarrays is not important. Use \(O(1)\) additional space and \(O(n)\) time.

This is almost exactly the same as the original partition problem where we can consider the 3 keys to be ` < p, = p , > p`.

We use the exact same algorithm skeleton and make a few changes as below:

def partition_keys(keys, A):
  """We define subarrays as follows:
  A[:k1]: elements with key 1
  A[k1:k2]: elements with key 2
  A[k2:k3]: unclassified elements
  A[k3:]: elements with key 3"""
  k1, k2, k3 = (0, 0, len(A)-1)
  # We use k2 as an iterator
  while k2 <= k3:
      if A[k2].key == keys['Key1']:
          A[k1], A[k2] = A[k2], A[k1]
          k2 += 1
          k1 += 1
      elif A[k2].key == keys['Key2']:
          k2 += 1
      elif A[k2].key == keys['Key3']:
          A[k2], A[k3] = A[k3], A[k2]
          k3 -= 1
  return None

Given an array of \(n\) objects with keys that take one of 4 values. Reorder the array so that all objects with same keys appear together. \(O(1)\) additional space and \(O(n)\) time.

Reusing the same algorithm as above, we maintain the following invariants with each swap:
1. A[:k1] = elements with key 1
2. A[k1:k2] = elements with key 2
3. A[k2:k3] = elements with key 3
4. A[k3:k4] = ? unclassified elements
5. A[k4:] = elements with key 4

def partition_4_keys(keys, A):
    """We define invariants as follows:
    A[:k1]: elements with key 1
    A[k1:k2]: elements with key 2
    A[k2:k3]: elements with key 3
    A[k3:k4]: unclassified elements
    A[k4:]: elements with key 4"""
    k1, k2, k3, k4 = (0, 0, 0, len(A) - 1)
    # We use k3 as an iterator
    while k3 >= k4:
        if A[k3].key == keys['Key1']:
            A[k1], A[k3] = A[k3], A[k1]
            k1 += 1
            k2 += 1
            k3 += 1
        elif A[k3].key == keys['Key2']:
            A[k2], A[k3] = A[k3], A[k2]
            k2 += 1
            k3 += 1
        elif A[k3].key == keys['Key3']:
            k3 += 1
        elif A[k3].key == keys['Key4']:
            A[k3], A[k4] = A[k4], A[k3]
            k4 -= 1
    return None

We can see a sort of pattern emerging if we wanted to reorder an array into k sub arrays then we would need to maintain k+1 invariants, and when checking for k1 we would update all key index boundaries k1,k2,k3 ..k-1. We would also use the k-1 key index as an iterator i.e: k3 is like an iterator for 4 keys, k2 is like an iterator for 3 keys. k5 would be an iterator for 6 keys… and so on.

Given an array of \(n\) objects with Boolean-valued keys, reorder the array so that objects with false keys appear first. \(O(1)\) space and \(O(n)\) time.
Given an array of \(n\) objects with Boolean-valued keys. Objects with false keys should appear first and the relative ordering of objects with key true should not change. \(O(1)\) space and \(O(n)\) time.

Twitter Facebook LinkedIn