MEDIUM
DATA STRUCTURES AND ALGORITHMS

How to Solve the Partition List Problem

Written By
kenny-polyack.png
Kenny Polyak
Jared Skinner

An Introduction to the Partition List Interview Question

The Partition List problem involves dividing an array into three parts based on a given value. The goal is to create three partitions: one containing all elements less than the given value, one with elements equal, and the other with elements greater. This problem requires carefully swapping array elements and a strong understanding of quicksort.

Examples of the List Partition Problem

Given a list of integers L and a number K, write a function that reorganizes L into three partitions: elements less than K, elements equal to K, and elements greater than K. No additional lists may be used.

** Note, relative ordering of elements need not be maintained.

Example Inputs and Outputs

Example 1

Input:
L = [6,1,2,4,5,3] K = 3 Output: L = [1,2,3,4,5,6] L = [2,1,3,6,4,5]

Example 2

Input: L = [1,2,3,4,5,6] K = 3 Output: L = [1,2,3,4,5,6] L = [2,1,3,6,4,5]

Example 3

Input: L = [7,0,6,4,10,7,3,-4,3,12] K = 7 Output: L = [0,6,4,3,-4,3,7,7,10,12] L = [-4,0,3,3,4,6,7,7,10,12]

Constraints

  • 1 <= n <= 10,000, where n is the number of elements in L

How to Solve the Partition List Problem (With Solutions in Python & JavaScript)

Below we'll walk through two different approaches to the partition list problem – brute force and applied quicksort – and provide walkthrough solutions in Python and JavaScript.

Approach 1: Brute Force

Our initial approach will be to iterate through all elements in the list and move elements less than K to the left side of the list and move elements greater than K to the right side of the list.

While this approach is simple enough on paper, there are some "gotcha's" we need to be aware of:

  • Any time we are making changes to a list while iterating over it we have to be especially careful. Moving elements to the beginning of the list or the end of the list could throw off our indexing. This could result in us iterating over an element twice or skipping an element altogether! Because of this, we need to make certain that each time we make a change to our list we assign our iterator the correct next index value.

  • Because we are moving elements to the end of our list, if we are not careful we could run into an endless loop where we are picking up the same elements over and over again and pushing them to the end of our list. To combat this we will count the number of outstanding elements we need to sort.

This solution uses no additional space, giving us a space complexity of O(1) which is great! With that said, our time complexity is nothing to write home about. We are moving a ton of elements to either the beginning or the end of our list. To do this our list has to be shifted back and forth to accommodate the moving element. Each shift is at worst an O(n) operation and we do this n times! This gives us a time complexity of O(n^2).

def solution(L, K):
  idx = 0
  elements_remaining = len(L)
 
  while elements_remaining > 0:
    if L[idx] == K:
      idx += 1
    elif L[idx] < K:
      val = L.pop(idx)
      L.insert(0, val)
      idx += 1
    else: # L[idx] > K       
      val = L.pop(idx)
      L.append(val)
      # don't update idx!
 
    elements_remaining -= 1
 
  return L

Time/Space Complexity

  • Time complexity: O(n^2)
  • Space complexity: O(1)

Approach 2: Applied Quicksort

Now, our space complexity is currently very good, but our time complexity leaves something to be desired. To address this we are going to look at a common tool when dealing with partitioning lists in order: Quicksort. Quicksort is a sorting algorithm used to efficiently sort a list in place (without using any extra lists). Quicksort works by identifying a "pivot" element, partitioning the list into elements less than the pivot and elements greater than the pivot, and then repeating the process on the left and right sides. This is done recursively until the entire list is sorted (see https://en.wikipedia.org/wiki/Quicksort for more details). While we don't need to sort the entire list, identifying a pivot and partitioning is exactly what we want!

In our first approach, we would move elements to the beginning or end of our list depending on how they compared to K. This is at worst an O(n) operation since each time we remove or add an element we have to shift the other elements in the list to accommodate. Instead of this add/remove approach, the efficiency of quicksort comes from swapping elements, as this is in itself a constant time operation. The problem is if we swap elements, we need TWO elements to swap. Up to now, we have been dealing with a single element at a time. Quicksort solves this issue by locating two elements in the "wrong place" and swapping them so we are one step closer to a partitioned list. Here's how the algorithm works:

  1. Select a pivot
  2. Swap our pivot with the first element in the list to get it out of the way.
  3. Walk through the elements of the list (excluding the pivot) from the outside in. Use two variables left_idx and right_idx to track which indices we are currently inspecting. Whenever we find two elements that are on the wrong sides of our partition, swap them.
  4. Swap the pivot with the last inspected left element

Regarding part 3 it can be confusing to think of walking through both ends of the array at the same time. Instead, think of walking through the left end of the array until you find a value greater than K and then walking through the right end of the array until you find a value less than k.

This approach is almost perfect for us. There's only one problem: we don't know how many times K will show up in our list! It could show up once, or ten times, or maybe not at all. The algorithm as described implicitly assumes that K is present in our list only once and can be used as a pivot. While this is an edge case that an interviewer will want to hear you consider, we need to consider the more general case as well. We will handle this with two cases:

Case 1. If K doesn't show up in the list, find the element whose value is closest to K and use this as a pivot.

Case 2. If K shows up in the list one or more times, move all copies of K to the beginning of the list. Just like with a single pivot, swap all copies of K into the middle of the list once partitioning is done.

def solution(L, K):
  # Select the element closest to K as our pivot
  # Any K's that are found should be shifted to the beginning of the array
  # We will keep track of these and sub them in later
  K_list = []
 
  pivot = None
  for idx, elem in enumerate(L):
    if elem == K:
      K_list.append(idx)
 
    if pivot is None:
      pivot = elem
      pivot_idx = idx
      continue
 
    if abs(pivot - K) > abs(elem - K):
      pivot = elem
      pivot_idx = idx
 
  # Move any Ks to the beginning of our list to keep
  # them out of the way during partitioning.
  K_count = 0
  if len(K_list) > 0:
    for idx in K_list:
      L[K_count], L[idx] = L[idx], L[K_count]
      K_count += 1
  else:
    # Swap our pivot to the beginning of our list
    L[pivot_idx], L[0] = L[0], L[pivot_idx]
 
  left_idx = K_count
  right_idx = len(L) - 1
  while left_idx < right_idx:
    while L[left_idx] < pivot and left_idx < len(L) - 1:
      left_idx += 1
 
    while L[right_idx] > pivot and right_idx > 0:
      right_idx -= 1
 
    # Make sure the left index is actually left of the right index!
    if left_idx < right_idx:
      L[left_idx], L[right_idx] = L[right_idx], L[left_idx]
 
  # Put the pivot (or Ks) back in place!
  for idx, val in enumerate(L):
    if val > pivot:
      break
 
    pivot_idx = idx
 
  offset = (pivot_idx + 1) - K_count
  while K_count > 0:
    i = K_count - 1
    L[i], L[i + offset] = L[i + offset], L[i]
    K_count -= 1
 
  return L

Time/Space Complexity

  • Time complexity: O(n)
  • Space complexity: O(1)

Practice the Partition List Problem With Our AI Interviewer

Start AI Interview

About interviewing.io

interviewing.io is a mock interview practice platform. We've hosted over 100K mock interviews, conducted by senior engineers from FAANG & other top companies. We've drawn on data from these interviews to bring you the best interview prep resource on the web.

We know exactly what to do and say to get the company, title, and salary you want.

Interview prep and job hunting are chaos and pain. We can help. Really.