Median of two sorted lists in Python

Why this matters

The median is the middle value of a sorted list. For a list of 5 values, the median is the third value. For a list of 4 values, there is no single middle value, so the median is the average of the two middle values. The median is often used to describe a typical value, because one very large number does not move it much.

A list is Python's numbered collection of values, written with square brackets, such as [1, 2, 3]. Each value has an index, which is its position. Python counts indexes from 0. The standard library function statistics.median returns the median of a list. It sorts the values first, so it works on unsorted lists too.

Given two sorted lists, the simple method merges them into one sorted list and reads the middle. Merging takes time proportional to the total size, which is m plus n when the lists have m and n values. The method in this lesson does not merge. It binary searches for the right place to cut the two lists. Its cost is O(log(min(m, n))), where min(m, n) is the size of the smaller list.

The binary search lessons earlier in this module search for a value. This lesson searches for a cut position instead. The Quickselect lesson later in the module covers finding the k-th smallest value, which is a different problem.

The median in small examples

Take the sorted list [1, 2, 3, 4, 5]. It has 5 values, so the median is 3.

Take the sorted list [1, 2, 3, 4]. It has 4 values. The two middle values are 2 and 3. The median is their average, (2 + 3) / 2, which is 2.5.

Now take two lists, [1, 3] and [2]. The simple method merges them into [1, 2, 3] and reads the middle value, 2. The answer is 2.0. This lesson's function always returns a float, which is Python's type for decimal numbers, so the answer prints with a decimal point. The binary search method gives the same answer.

The idea: cut both lists

Take A = [1, 3, 8, 9, 15] and B = [7, 11, 18, 19, 21, 25]. Together they hold 11 values. The median of 11 values is the 6th value in sorted order.

Imagine a cut in A and a cut in B. The values before each cut go on the left side. The values after each cut go on the right side. Choose the cuts so that the left side holds exactly 6 values, which is half of 11 rounded up.

The cut is correct when two things are true:

Every value on the left is at most every value on the right.
For an odd total, the largest value on the left is the median.

Both lists are sorted, so you only need to check two pairs of values:

The last value on the left of A is at most the first value on the right of B.
The last value on the left of B is at most the first value on the right of A.

Take this cut. A has 4 values on the left: [1, 3, 8, 9]. B has 2 values on the left: [7, 11]. That gives 6 values on the left.

The last left value of A is 9. The first right value of B is 18. Since 9 is at most 18, the first pair is fine.
The last left value of B is 11. The first right value of A is 15. Since 11 is at most 15, the second pair is fine.

The largest value on the left is 11, which is the median. The function prints 11.0 for this input.

Once you choose how many values A gives to the left side, the number from B is fixed. Left side size is 6, so B gives 6 minus the number from A. This means the search only has to choose one number: how many values to take from A.

Binary search on the cut in A

Let i be the number of values taken from A for the left side. Let j be the number taken from B. The code sets j to the left side size minus i.

Now compare the two checks.

If the last left value of A is greater than the first right value of B, then i is too big. Too many values from A are on the left. Move the search to smaller values of i.
If the last left value of B is greater than the first right value of A, then i is too small. Too few values from A are on the left. Move the search to larger values of i.
If neither check fails, the cut is correct.

The search range for i starts at 0 and ends at the size of A. Each step halves the range, so the number of steps is about the base 2 logarithm of the size of A (how many times A can be halved before one value is left).

This is why the code makes A the shorter list. If A is the longer one, the range for i is larger, and j can fall outside B. The code swaps the two lists at the start when A is longer, with a, b = b, a. This is a tuple swap. A tuple is a fixed group of values written with commas, such as b, a. Python builds the pair (b, a) first and then assigns both names, so no temporary variable is needed. Swapping does not change the median, because the median of two lists does not depend on which one you call first.

A trace of the example

The left side size is 6, so the search starts with lo at 0 and hi at 5.

Step 1 tries i = 2. That takes 2 values from A and 4 from B. The check on the first pair passes, because 3 is at most 21. The check on the second pair fails, because 19 is greater than 8. So too few values come from A. The search moves lo up to 3.

Step 2 tries i = 4. That takes 4 values from A and 2 from B. Both checks pass, so the cut is correct. The total is odd, so the median is the largest value on the left, which is 11.

Why binary search works here

The two checks change in a predictable way as i changes. When i goes up by one, A gives one more value to the left side, and B gives one fewer. So A's values on the left go up, and B's values on the right go down.

The first check is "A's last left value is greater than B's first right value". Once it turns true for some i, it stays true for every larger i. The second check is "B's last left value is greater than A's first right value", and it works the other way round. It stays true for every smaller i.

Because of this, the search can always tell which way to move. A binary search only needs that one rule to be true for the whole range.

A second example with the cut in the other list

Take A = [10, 20, 30] and B = [5, 15, 25, 35, 45]. Together they hold 8 values, so the left side needs 4 values. The median is the average of the largest left value and the smallest right value.

The trace below shows the search.

step 1: lo=0 hi=3 i=1 j=3 aLeft=10 aRight=20 bLeft=25 bRight=35
step 2: lo=2 hi=3 i=2 j=2 aLeft=20 aRight=30 bLeft=15 bRight=25
found partition after 2 steps

Step 1 takes 1 value from A and 3 from B. The first check passes, since 10 is at most 35. The second check fails, since 25 is greater than 20. So the search moves right.

Step 2 takes 2 values from A and 2 from B. Both checks pass. The largest left value is 20, and the smallest right value is 25. The total is even, so the median is (20 + 25) / 2, which is 22.5.

A third example with repeated values

Take A = [1, 1, 1] and B = [1, 1, 1, 1]. Every value is 1, so the median is 1.0.

step 1: lo=0 hi=3 i=1 j=3 aLeft=1 aRight=1 bLeft=1 bRight=1
found partition after 1 steps

The checks use <=, so equal values count as correctly ordered, and the first try passes. The common mistakes section shows what happens when the checks use < instead.

The code

The function find_median does the search. Python has no fixed-size integer type that could overflow here, so the code only needs sentinels, which are stand-in values, for an empty side. INF = float("inf") is Python's positive infinity, and -INF is negative infinity. Every number is smaller than INF and larger than -INF. A value of -INF on the left never breaks a check, and INF on the right never breaks a check either. If you know Java, these two values play the role of Integer.MIN_VALUE and Integer.MAX_VALUE.

The expression -INF if i == 0 else a[i - 1] is a conditional expression. It reads as: if i is 0, use -INF, otherwise use a[i - 1].

Python · editable

INF = float("inf")

def find_median(a, b):
    """Return the median of two sorted lists a and b."""
    if len(a) > len(b):
        a, b = b, a  # make a the shorter list
    m = len(a)
    n = len(b)
    if m + n == 0:
        raise ValueError("both lists are empty")

half = (m + n + 1) // 2  # how many values the left side must hold
    lo = 0
    hi = m
    while lo <= hi:
        i = lo + (hi - lo) // 2  # values taken from a
        j = half - i             # values taken from b

a_left = -INF if i == 0 else a[i - 1]
        a_right = INF if i == m else a[i]
        b_left = -INF if j == 0 else b[j - 1]
        b_right = INF if j == n else b[j]

if a_left <= b_right and b_left <= a_right:
            max_left = max(a_left, b_left)
            if (m + n) % 2 == 1:
                return float(max_left)
            min_right = min(a_right, b_right)
            return (max_left + min_right) / 2
        elif a_left > b_right:
            hi = i - 1
        else:
            lo = i + 1
    raise ValueError("input is not sorted")

examples = [
    ([1, 3], [2]),
    ([1, 2], [3, 4]),
    ([1, 3, 8, 9, 15], [7, 11, 18, 19, 21, 25]),
    ([], [4]),
    ([5], []),
]
for a, b in examples:
    print(f"{a} and {b} -> median {find_median(a, b)}")

The output is:

[1, 3] and [2] -> median 2.0
[1, 2] and [3, 4] -> median 2.5
[1, 3, 8, 9, 15] and [7, 11, 18, 19, 21, 25] -> median 11.0
[] and [4] -> median 4.0
[5] and [] -> median 5.0

The value half is the number of values the left side must hold. Using (m + n + 1) // 2 puts the extra value on the left when the total is odd. That is why, for an odd total, the answer is the largest value on the left.

Take [1, 2] and [3, 4], which total 4 values. The left side must hold 2 values. The first try is i = 1, which takes 1 value from each list. The first check passes, because A's last left value 1 is at most B's first right value 4. The second check fails, because B's last left value 3 is greater than A's first right value 2. So the search moves right. The next try is i = 2, which takes both values from A. The left side is [1, 2], and the right side is [3, 4]. The largest left value is 2, and the smallest right value is 3. The even total gives (2 + 3) / 2, which is 2.5.

The case [] and [4] shows the sentinels. A has no values, so i is always 0. The left side gets the value 4 from B, and the answer is 4.0.

Why the search takes so few steps

The search range for i has about m + 1 values. Each step halves that range. For a list of 1,000 values, about 10 steps are enough, because 2 to the power of 10 is 1,024. The next section counts the steps for a pair of lists with 1,000 and 1,000,000 values. The second list is much bigger, but it only changes the work of each step, not the number of steps.

Showing each step

The program below copies the search and prints each step for the example above. It also counts the steps for the large pair.

Python · editable

INF = float("inf")

def trace(a, b):
    m = len(a)
    n = len(b)
    half = (m + n + 1) // 2
    lo = 0
    hi = m
    step = 0
    while lo <= hi:
        i = lo + (hi - lo) // 2
        j = half - i
        a_left = -INF if i == 0 else a[i - 1]
        a_right = INF if i == m else a[i]
        b_left = -INF if j == 0 else b[j - 1]
        b_right = INF if j == n else b[j]
        step += 1
        print(f"step {step}: lo={lo} hi={hi} i={i} j={j} "
              f"aLeft={a_left} aRight={a_right} bLeft={b_left} bRight={b_right}")
        if a_left <= b_right and b_left <= a_right:
            print(f"found partition after {step} steps")
            return
        elif a_left > b_right:
            hi = i - 1
        else:
            lo = i + 1

def steps(m, n):
    a = [2 * i for i in range(m)]
    b = [2 * i + 1 for i in range(n)]
    half = (m + n + 1) // 2
    lo = 0
    hi = m
    count = 0
    while lo <= hi:
        count += 1
        i = lo + (hi - lo) // 2
        j = half - i
        a_left = -INF if i == 0 else a[i - 1]
        a_right = INF if i == m else a[i]
        b_left = -INF if j == 0 else b[j - 1]
        b_right = INF if j == n else b[j]
        if a_left <= b_right and b_left <= a_right:
            return count
        elif a_left > b_right:
            hi = i - 1
        else:
            lo = i + 1
    return -1

trace([1, 3, 8, 9, 15], [7, 11, 18, 19, 21, 25])
print("steps for 1000 and 1000000:", steps(1000, 1000000))

The output is:

step 1: lo=0 hi=5 i=2 j=4 aLeft=3 aRight=8 bLeft=19 bRight=21
step 2: lo=3 hi=5 i=4 j=2 aLeft=9 aRight=15 bLeft=11 bRight=18
found partition after 2 steps
steps for 1000 and 1000000: 10

The 10 steps for 1,000 and 1,000,000 values match the estimate. The search runs over the 1,000 values of the smaller list, not over the million.

Common mistakes

Reading past the start or end of a list

Python does not treat a negative index as an error. A negative index counts from the end of the list, so a[-1] is the last value. If the left side of A is empty, i is 0, and a[i - 1] is a[-1]. The program does not stop. It uses the wrong value and carries on.

Python · editable

The output is 15, which is the last value of A, not the value on the left of the cut.

An index past the end does stop the program. Calling left_no_check(a, 6) reads a[5], and the list only has indexes 0 to 4. The last line of the error is:

IndexError: list index out of range

The fix is the i == 0 check in the code above. The sentinel value is used instead of reading the list.

Dividing a very large whole number with `/`

Python integers have no fixed size, so adding two large values does not wrap around to a negative number:

Python · editable

The output is:

4000000000
2000000000.0

The / operator always returns a float. A float keeps only about 53 bits of precision, so a whole number above 2 to the power of 53 can lose its last digits. The // operator keeps the result as an exact integer.

Python · editable

The output is:

exact with //: 9007199254740993
with /: 9007199254740992.0
big with //: 1152921504606846977
big with /: 1.152921504606847e+18

A float also has a largest size. Dividing a number that is too large for a float raises an error. For example, (10 ** 400 + 10 ** 400) / 2 stops with this line at the end of the error:

OverflowError: integer division result too large for a float

Use // when the values are whole numbers and the answer must be exact. The function in this lesson uses / only in the even case, where the answer is allowed to be a float.

Using the wrong size for the left side

If you write (m + n) // 2 instead of (m + n + 1) // 2, the left side holds one value fewer when the total is odd. Then the largest value on the left can be the value just below the median. The + 1 puts the extra value on the left, so that the largest left value is the median.

Searching the longer list

If you do not swap the lists, the search range can be larger than the shorter list. Then j can go below 0 or above the size of B. If j goes below 0, Python reads from the end of B without an error, as shown above. If j goes past the end, the program raises an IndexError. The swap at the top of find_median keeps j inside B.

Using `<` instead of `<=`

With duplicate values, a correct cut can have equal values on both sides of it. If the checks use < instead of <=, a correct cut never passes. In a copy of find_median named find_median_strict, the two checks become:

if a_left < b_right and b_left < a_right:

With that change, find_median_strict([1, 1, 1], [1, 1, 1, 1]) ends the loop without a match, and the function raises an error. The last line of the error is:

ValueError: input is not sorted

The message blames the input, but the real problem is the comparison.

Passing unsorted input

The function assumes both lists are sorted. It cannot check this fully. Unsorted input can give a wrong answer without any error. Calling find_median([3, 1, 2], [5]) returns 1.5. The merged sorted list is [1, 2, 3, 5], so the median is 2.5. Sort each list before you call the function.

Practice

Exercise 1

What is the median of [2, 4, 6] and [1, 3, 5, 7]?

Show the answer

The merged list is [1, 2, 3, 4, 5, 6, 7]. It has 7 values, so the median is the 4th value, 4. The function returns 4.0.

Exercise 2

What is the median of [10, 20] and [30, 40]?

Show the answer

The merged list is [10, 20, 30, 40]. The two middle values are 20 and 30, so the median is (20 + 30) / 2, which is 25. The function returns 25.0.

Exercise 3

What does find_median return for [] and [1, 2, 3]?

Show the answer

The median of [1, 2, 3] is 2. The function returns 2.0. In this case the empty list is the shorter one, so it is already a, and the search has only one possible value of i, which is 0.

Exercise 4

Why does the code swap the lists so that a is always the shorter one?

Show the answer

The search range for i runs from 0 to the length of a. The left side must hold half values in total, so j equals half minus i. If a were longer than b, then j could be larger than the length of b, and the code would read past the end of b. Keeping a shorter keeps j inside b.

Exercise 5

What is the median of [10, 30, 50] and [20, 40]? Write the merged list first.

Show the answer

The merged list is [10, 20, 30, 40, 50]. It has 5 values, so the median is the 3rd value, 30. The function returns 30.0.

Exercise 6

For [1, 2] and [3, 4], how many steps does the search take? Give the value of i on each step.

Show the answer

There are 2 steps. The first step tries i = 1, and the second check fails, so the search moves to larger values. The second step tries i = 2, and both checks pass. The answer is 2.5.

Summary

The median of two sorted lists can be found without merging them.
The search picks a cut in the shorter list. The cut in the other list is then fixed by the left side's size.
Two comparisons decide whether a cut is correct, and they decide which way the search moves.
The search takes about log2 of the shorter length in steps, which is much less than the total length.
Use float("inf") for empty sides, and use // when the values are whole numbers and the answer must be exact.

Median of two sorted lists in Python

Why this matters

The median in small examples

The idea: cut both lists

Binary search on the cut in A

A trace of the example

Why binary search works here

A second example with the cut in the other list

A third example with repeated values

The code

Why the search takes so few steps

Showing each step

Common mistakes

Reading past the start or end of a list

Dividing a very large whole number with /

Using the wrong size for the left side

Searching the longer list

Using < instead of <=

Passing unsorted input

Practice

Exercise 1

Exercise 2

Exercise 3

Exercise 4

Exercise 5

Exercise 6

Summary

Stuck? Ask the tutor

Dividing a very large whole number with `/`

Using `<` instead of `<=`