Sorting, Intervals and Quickselect

Sorting is rarely the question, and often the first step of the answer. Once data is ordered, many problems collapse: neighbours become comparable, duplicates become adjacent, and intervals line up. This chapter covers what to know about sorting algorithms, how to use sorting as a tool, the interval problems that depend on it, and quickselect for finding the $k$th element without a full sort.

1. What you should know about sorting algorithms

You will usually call the built-in sort. You should still be able to explain the main algorithms and their trade-offs.

Two facts matter in practice:

• Comparison sorts cannot beat $O(n\log n)$ in the worst case. This is why "sort the array" costs $O(n\log n)$ in your analysis.
• A sort is stable if equal elements keep their original order. Python's sort and sorted are stable, which lets you sort by several keys in passes, or use a tuple as the key.

Python's built-in sort (Timsort) is $O(n\log n)$ worst case and exploits existing runs, so it is very fast on partly ordered data.

people = [("ann", 30), ("bob", 25), ("cat", 30), ("dan", 25)]
people.sort(key=lambda p: (-p[1], p[0]))        # age descending, then name ascending
assert people == [("ann", 30), ("cat", 30), ("bob", 25), ("dan", 25)]

Merge sort and quicksort, briefly

Merge sort splits the list in half, sorts each half and merges (shown in the complexity chapter). Quicksort picks a pivot, partitions the list into smaller, equal and larger parts, and recurses on the parts. A random pivot makes the quadratic worst case extremely unlikely.

import random

def quicksort(a):
    if len(a) <= 1:
        return a
    pivot = random.choice(a)
    smaller = [x for x in a if x < pivot]
    equal = [x for x in a if x == pivot]
    larger = [x for x in a if x > pivot]
    return quicksort(smaller) + equal + quicksort(larger)

assert quicksort([3, 6, 1, 8, 2, 9, 2]) == [1, 2, 2, 3, 6, 8, 9]
assert quicksort([]) == []

This version is clear but uses extra memory. The in-place partition scheme is a favourite follow-up, and the quickselect code later in this chapter shows it.

2. Counting sort and bucket ideas

When keys are integers in a small range, count them instead of comparing.

def counting_sort(nums, max_value):
    counts = [0] * (max_value + 1)
    for x in nums:
        counts[x] += 1
    out = []
    for value, c in enumerate(counts):
        out.extend([value] * c)
    return out                                  # O(n + max_value)

assert counting_sort([4, 2, 2, 8, 3, 3, 1], 8) == [1, 2, 2, 3, 3, 4, 8]

Bucket thinking also solves "top $k$ frequent" in $O(n)$ (arrays chapter) and the maximum gap problem.

3. Sorting as a preprocessing step

A large family of problems is: sort first, then do something simple.

Meeting rooms. Can one person attend all meetings? Sort by start and check that each begins after the previous one ends.

def can_attend_all(intervals):
    intervals.sort(key=lambda x: x[0])
    for i in range(1, len(intervals)):
        if intervals[i][0] < intervals[i - 1][1]:        # overlaps the previous meeting
            return False
    return True

assert can_attend_all([[0, 30], [5, 10], [15, 20]]) is False
assert can_attend_all([[7, 10], [2, 4]]) is True

Largest number. Arrange numbers to form the largest concatenation. Sort with a custom comparison: a before b if a + b > b + a as strings.

from functools import cmp_to_key

def largest_number(nums):
    strs = [str(n) for n in nums]
    strs.sort(key=cmp_to_key(lambda a, b: -1 if a + b > b + a else (1 if a + b < b + a else 0)))
    result = "".join(strs)
    return "0" if result[0] == "0" else result

assert largest_number([10, 2]) == "210"
assert largest_number([3, 30, 34, 5, 9]) == "9534330"
assert largest_number([0, 0]) == "0"

The lesson: sorting with a custom ordering is powerful when you can prove the ordering is consistent (transitive).

4. Intervals

Interval problems give ranges [start, end] and ask about overlaps, merges and counts. The universal first move is sort by start time. After sorting, overlapping intervals are neighbours, and you can sweep once.

Merge intervals

def merge_intervals(intervals):
    intervals.sort(key=lambda x: x[0])
    merged = []
    for start, end in intervals:
        if merged and start <= merged[-1][1]:          # overlaps the last merged interval
            merged[-1][1] = max(merged[-1][1], end)
        else:
            merged.append([start, end])
    return merged

assert merge_intervals([[1, 3], [2, 6], [8, 10], [15, 18]]) == [[1, 6], [8, 10], [15, 18]]
assert merge_intervals([[1, 4], [4, 5]]) == [[1, 5]]
assert merge_intervals([[1, 4], [2, 3]]) == [[1, 4]]
assert merge_intervals([]) == []

Time $O(n\log n)$ for the sort, then $O(n)$. Decide what "overlap" means for touching intervals ([1,4] and [4,5]): here they merge, because start <= last_end. Ask the interviewer, since the answer changes the comparison from <= to <.

Insert interval

Insert a new interval into a sorted, non-overlapping list, merging if needed. Three phases: add all intervals that end before the new one starts, merge all that overlap it, add the rest.

def insert_interval(intervals, new):
    result = []
    i, n = 0, len(intervals)
    while i < n and intervals[i][1] < new[0]:          # entirely before
        result.append(intervals[i]); i += 1
    while i < n and intervals[i][0] <= new[1]:         # overlapping: absorb
        new = [min(new[0], intervals[i][0]), max(new[1], intervals[i][1])]
        i += 1
    result.append(new)
    result.extend(intervals[i:])                       # entirely after
    return result

assert insert_interval([[1, 3], [6, 9]], [2, 5]) == [[1, 5], [6, 9]]
assert insert_interval([[1, 2], [3, 5], [6, 7], [8, 10], [12, 16]], [4, 8]) == [[1, 2], [3, 10], [12, 16]]
assert insert_interval([], [5, 7]) == [[5, 7]]

Time $O(n)$, because the input is already sorted.

Minimum number of meeting rooms (heap or sweep)

How many rooms are needed for a set of meetings? Equivalent: the maximum number of overlapping intervals at any moment.

Sweep line with sorted events. Separate all start times and all end times, sort each, and walk through them with two pointers: a start before the earliest end means a new room is needed, otherwise a room frees up.

def min_meeting_rooms(intervals):
    starts = sorted(i[0] for i in intervals)
    ends = sorted(i[1] for i in intervals)
    rooms = best = 0
    e = 0
    for s in starts:
        if s < ends[e]:
            rooms += 1                 # a meeting starts before the earliest one ends
        else:
            e += 1                     # a room is reused
        best = max(best, rooms)
    return best

assert min_meeting_rooms([[0, 30], [5, 10], [15, 20]]) == 2
assert min_meeting_rooms([[7, 10], [2, 4]]) == 1
assert min_meeting_rooms([[1, 5], [2, 6], [3, 7]]) == 3

Heap version. Keep a min-heap of end times of rooms in use. For each meeting by start time, if the earliest end is not after the start, reuse that room (pop), then push this meeting's end. The heap size at the end is the number of rooms.

import heapq

def min_meeting_rooms_heap(intervals):
    heap = []                                   # end times of meetings in progress
    for start, end in sorted(intervals):
        if heap and heap[0] <= start:
            heapq.heapreplace(heap, end)        # reuse the freed room
        else:
            heapq.heappush(heap, end)
    return len(heap)

assert min_meeting_rooms_heap([[0, 30], [5, 10], [15, 20]]) == 2
assert min_meeting_rooms_heap([[1, 5], [2, 6], [3, 7]]) == 3

Non-overlapping intervals (greedy)

Remove the fewest intervals so the rest do not overlap. Sort by end time and greedily keep the interval that ends earliest, because it leaves the most room. Count the ones you must drop.

def erase_overlap_intervals(intervals):
    intervals.sort(key=lambda x: x[1])
    kept_end = float("-inf")
    removed = 0
    for start, end in intervals:
        if start >= kept_end:
            kept_end = end                    # keep this one
        else:
            removed += 1                      # overlaps the one we kept: drop it
    return removed

assert erase_overlap_intervals([[1, 2], [2, 3], [3, 4], [1, 3]]) == 1
assert erase_overlap_intervals([[1, 2], [1, 2], [1, 2]]) == 2
assert erase_overlap_intervals([[1, 2], [2, 3]]) == 0

Note the sort key differs between problems: sort by start to merge, by end to select the maximum number of non-overlapping intervals. Knowing why is the point: the earliest finishing interval is always safe to keep (an exchange argument, see the greedy chapter).

5. Quickselect: the $k$th element in $O(n)$ on average

To find the $k$th largest element, sorting costs $O(n\log n)$ and a heap costs $O(n\log k)$. Quickselect partitions like quicksort but recurses into only the side that contains the target, giving $O(n)$ on average.

import random

def kth_largest(nums, k):
    target = len(nums) - k                     # index of the answer in sorted order

    def partition(lo, hi):
        pivot_index = random.randint(lo, hi)
        nums[pivot_index], nums[hi] = nums[hi], nums[pivot_index]
        pivot = nums[hi]
        store = lo
        for i in range(lo, hi):
            if nums[i] < pivot:
                nums[store], nums[i] = nums[i], nums[store]
                store += 1
        nums[store], nums[hi] = nums[hi], nums[store]
        return store

    lo, hi = 0, len(nums) - 1
    while lo <= hi:
        p = partition(lo, hi)
        if p == target:
            return nums[p]
        if p < target:
            lo = p + 1
        else:
            hi = p - 1

assert kth_largest([3, 2, 1, 5, 6, 4], 2) == 5
assert kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4) == 4
assert kth_largest([1], 1) == 1

Time $O(n)$ average ($n+n/2+n/4+\dots$), $O(n^2)$ worst case, made very unlikely by a random pivot. Space $O(1)$. It modifies the array in place. The decision between the heap and quickselect is a typical interview discussion: the heap is simpler and works on streams, quickselect is faster on a static array.

6. Merge sort variations

Merge sort's merge step can count extra information while merging.

Count of smaller numbers after self / count inversions. During the merge, when an element from the right half is placed before elements of the left half, those remaining left elements form inversions with it.

def count_inversions(nums):
    def sort(a):
        if len(a) <= 1:
            return a, 0
        mid = len(a) // 2
        left, inv_l = sort(a[:mid])
        right, inv_r = sort(a[mid:])
        merged, i, j, inv = [], 0, 0, inv_l + inv_r
        while i < len(left) and j < len(right):
            if left[i] <= right[j]:
                merged.append(left[i]); i += 1
            else:
                merged.append(right[j]); j += 1
                inv += len(left) - i            # all remaining left elements exceed right[j]
        merged.extend(left[i:]); merged.extend(right[j:])
        return merged, inv
    return sort(nums)[1]

assert count_inversions([2, 4, 1, 3, 5]) == 3
assert count_inversions([5, 4, 3, 2, 1]) == 10
assert count_inversions([1, 2, 3]) == 0

7. Choosing among the techniques

8. Common mistakes

• Sorting by the wrong key in an interval problem (start versus end).
• Mishandling touching intervals (< versus <=). Ask what counts as overlapping.
• Mutating the input by sorting it in place when the caller does not expect that. Use sorted(...) if needed.
• Forgetting that sorting changes indexes, when the problem asks for original positions.
• Using quickselect without a random pivot on sorted input, which hits the quadratic worst case.
• Assuming sort is $O(n)$ in complexity analysis. It is $O(n\log n)$.
• Inconsistent comparators, which can break the sort or give wrong results.

9. Practice set

1. Merge intervals, insert interval, interval list intersections.
2. Meeting rooms I and II, minimum number of arrows to burst balloons.
3. Non-overlapping intervals, car pooling (sweep line with a difference array).
4. Sort colours, sort an array by parity, relative sort array.
5. Kth largest element, top $k$ frequent, $k$ closest points to the origin.
6. Largest number, custom sort string.
7. Count inversions, count of smaller numbers after self.
8. Maximum gap (bucket idea), H-index.
9. Merge sorted array in place from the back.
10. Employee free time (merge many sorted interval lists).