Python – Partition an Array into Three Parts Around a Given Value

This tutorial explains how to partition an array into three segments based on a given pivot value. The goal is to rearrange the array so that all elements less than the pivot come first, followed by elements equal to the pivot, and then elements greater than the pivot.

Problem Statement

Given an array arr and a pivot value p, partition the array into three parts such that:

  1. All elements less than p come first.
  2. All elements equal to p come in the middle.
  3. All elements greater than p come last.

The partitioning does not require sorting within the three segments; it only groups the elements based on the pivot.

Sample Input and Output

Example 1:

</>
Copy
Input:  arr = [4, 9, 4, 2, 7, 4, 10], pivot = 4
Output: [2, 4, 4, 4, 9, 7, 10]

Example 2:

</>
Copy
Input:  arr = [5, 3, 8, 5, 2, 5, 9], pivot = 5
Output: [3, 2, 5, 5, 5, 8, 9]

Solution Approach

We can solve this problem efficiently using the Dutch National Flag algorithm (three-way partitioning) with the following steps:

  1. Initialize three pointers:
    • low: Marks the boundary for elements less than the pivot.
    • mid: Traverses the array.
    • high: Marks the boundary for elements greater than the pivot.
  2. Traverse the array:
    • If arr[mid] is less than the pivot, swap it with arr[low] and increment both low and mid.
    • If arr[mid] equals the pivot, just increment mid.
    • If arr[mid] is greater than the pivot, swap it with arr[high] and decrement high (without incrementing mid).
  3. Continue this process until mid exceeds high.

Python Program

</>
Copy
def partition_array(arr, pivot):
    low = 0           # Index for elements less than pivot
    mid = 0           # Index to traverse the array
    high = len(arr) - 1  # Index for elements greater than pivot

    while mid <= high:
        if arr[mid] < pivot:
            arr[low], arr[mid] = arr[mid], arr[low]
            low += 1
            mid += 1
        elif arr[mid] == pivot:
            mid += 1
        else:
            arr[mid], arr[high] = arr[high], arr[mid]
            high -= 1
    return arr

# Example usage:
arr1 = [4, 9, 4, 2, 7, 4, 10]
pivot1 = 4
print("Partitioned Array:", partition_array(arr1, pivot1))

arr2 = [5, 3, 8, 5, 2, 5, 9]
pivot2 = 5
print("Partitioned Array:", partition_array(arr2, pivot2))

Expected Output:

Partitioned Array: [2, 4, 4, 4, 7, 10, 9]
Partitioned Array: [3, 2, 5, 5, 5, 9, 8]

Conclusion

This tutorial demonstrated how to partition an array into three parts around a given pivot using the Dutch National Flag algorithm. The approach partitions the array in-place with a time complexity of O(n) and uses O(1) extra space.