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.
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.
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]
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]
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 LOur 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
1def solution(L, K):
2 idx = 0
3 elements_remaining = len(L)
4
5 while elements_remaining > 0:
6 if L[idx] == K:
7 idx += 1
8 elif L[idx] < K:
9 val = L.pop(idx)
10 L.insert(0, val)
11 idx += 1
12 else: # L[idx] > K
13 val = L.pop(idx)
14 L.append(val)
15 # don't update idx!
16
17 elements_remaining -= 1
18
19 return L
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:
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
1def solution(L, K):
2 # Select the element closest to K as our pivot
3 # Any K's that are found should be shifted to the beginning of the array
4 # We will keep track of these and sub them in later
5 K_list = []
6
7 pivot = None
8 for idx, elem in enumerate(L):
9 if elem == K:
10 K_list.append(idx)
11
12 if pivot is None:
13 pivot = elem
14 pivot_idx = idx
15 continue
16
17 if abs(pivot - K) > abs(elem - K):
18 pivot = elem
19 pivot_idx = idx
20
21 # Move any Ks to the beginning of our list to keep
22 # them out of the way during partitioning.
23 K_count = 0
24 if len(K_list) > 0:
25 for idx in K_list:
26 L[K_count], L[idx] = L[idx], L[K_count]
27 K_count += 1
28 else:
29 # Swap our pivot to the beginning of our list
30 L[pivot_idx], L[0] = L[0], L[pivot_idx]
31
32 left_idx = K_count
33 right_idx = len(L) - 1
34 while left_idx < right_idx:
35 while L[left_idx] < pivot and left_idx < len(L) - 1:
36 left_idx += 1
37
38 while L[right_idx] > pivot and right_idx > 0:
39 right_idx -= 1
40
41 # Make sure the left index is actually left of the right index!
42 if left_idx < right_idx:
43 L[left_idx], L[right_idx] = L[right_idx], L[left_idx]
44
45 # Put the pivot (or Ks) back in place!
46 for idx, val in enumerate(L):
47 if val > pivot:
48 break
49
50 pivot_idx = idx
51
52 offset = (pivot_idx + 1) - K_count
53 while K_count > 0:
54 i = K_count - 1
55 L[i], L[i + offset] = L[i + offset], L[i]
56 K_count -= 1
57
58 return L
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.
Interview prep and job hunting are chaos and pain. We can help. Really.