（转载）我最喜欢的算法：线性时间查找中位数

原文链接：My Favorite Algorithm: Linear Time Median Finding

在一个数组中寻找中位数似乎是一件小事，但是如果要在线性时间内做到就比较棘手。在这篇文章中，我将介绍我最喜欢的算法之一，the median-of-medians 方法，它可以在线性时间内找到数组中的中位数。虽然证明这个算法是线性的有点棘手，但是只有基础算法分析的读者也可以看懂。

时间复杂度为O(n log n)的中位数查找算法

找到中位数最直接的方法是对数组进行排序，然后可以通过索引得到中位数。快排的复杂度为$O(nlogn)$，这占了大部分时间。

def nlogn_median(l):
	l = sorted(l)
	if len(l)%2 == 1:
		return l[len(l) / 2]
	else:
		return 05 * (l[len(l)/2 - 1] + )

虽然这个方法的代码最简单，但它肯定不是最快的。

平均时间复杂度为O(n)的中位数查找算法

这个算法被称为快速选择，由Tony Hoare开发的，它也发明了快速排序算法。它是一个递归算法可以找到任何元素(不仅仅是中位数)。

1. 选择数组的一个索引，如何选择并不重要。但实践中随机选择一个效果很好。这个索引处的元素称为pivot。
2. 将数组一分为二：
   • 小于或等于pivot的元素：lesser_els
   • 大于pivot的元素：great_els
3. 中位数肯定在lesser_els和great_else中的一个。如果我们正在寻找第K个元素。
   • 如果lesser_els中有大于等于K个元素，则在lesser_els中递归搜索第K个元素。
   • 如果lesser_els中的元素个数小于K，则在greater_els中递归。不是搜索K，而是搜索K-len(lesser_els)。

下面是一个示例：

Consider the list below. We'd like to find the median.
l = [9,1,0,2,3,4,6,8,7,10,5]
len(l) == 11, so we're looking for the 6th smallest element
First, we must pick a pivot. We randomly select index 3.
The value at this index is 2.

Partitioning based on the pivot:
[1,0,2], [9,3,4,6,8,7,10,5]
We want the 6th element. 6-len(left) = 3, so we want
the third smallest element in the right array

We're now looking for third smallest element in the array below:
[9,3,4,6,8,7,10,5]
We pick an index at random to be our pivot.
We pick index 3, the value at which, l[3]=6

Partitioning based on the pivot:
[3,4,5,6] [9,7,10]
We want the 3rd smallest element, so we know it's the
3rd smallest element in the left array

We're now looking for the 3rd smallest in the array below:
[3,4,5,6]
We pick an index at random to be our pivot.
We pick index 1, the value at which, l[1]=4
Partitioning based on the pivot:
[3,4] [5,6]
We're looking for the item at index 3, so we know it's
the smallest in the right array.

We're now looking for the smallest element in the array below:
[5,6]

At this point, we can have a base case that chooses the larger
or smaller item based on the index.
We're looking for the smallest item, which is 5.
return 5

代码如下：

import random
def quickselect_median(l, pivot_fn=random.choice):
    if len(l) % 2 == 1:
        return quickselect(l, len(l) // 2, pivot_fn)
    else:
        return 0.5 * (quickselect(l, len(l) / 2 - 1, pivot_fn) +
                      quickselect(l, len(l) / 2, pivot_fn))


def quickselect(l, k, pivot_fn):
    """
    Select the kth element in l (0 based)
    :param l: List of numerics
    :param k: Index
    :param pivot_fn: Function to choose a pivot, defaults to random.choice
    :return: The kth element of l
    """
    if len(l) == 1:
        assert k == 0
        return l[0]

    pivot = pivot_fn(l)

    lows = [el for el in l if el < pivot]
    highs = [el for el in l if el > pivot]
    pivots = [el for el in l if el == pivot]

    if k < len(lows):
        return quickselect(lows, k, pivot_fn)
    elif k < len(lows) + len(pivots):
        # We got lucky and guessed the median
        return pivots[0]
    else:
        return quickselect(highs, k - len(lows) - len(pivots), pivot_fn)

快速选择几乎没有运行开销，平均时间复杂度为$O(n)$，接下来做以证明。

平均O(n)的证明

平均清况下，pivot会将数组划分为几乎同等元素个数大小的2部分。因此，每个后续的递归都只对上一步1/2的数据进行操作。

$$C=n+\frac{n}{2}+\frac{n}{4}+\frac{n}{8}+\cdots =2n=O(n)$$

有很多方法可以证明这个连加和趋近于$2n$。可以参考这篇文章。

快速选择在理想情况下(上面所示)是可以达到$O(n)$复杂度的，但如果划分时不平均，但同时想保证我们的算法无论如何都是线性时间，该怎么办？

确定O(n)的证明

在上文中阐述了平均性能为$O(n)$的快速选择算法。从技术上来讲，如果每一步选择的pivot都是当前数组中的最大元素，则每个步骤只会从数组中删除一个元素，实际的复杂度为$O(n^2)$。所以pivot选择的好坏直接影响时间复杂度。

考虑到这一点，接下来介绍一个用于选择pivot的算法。我们的目的是在线性时间里选择一个pivot提供给快速选择算法使用。这个算法最初是由Blum, Floyd, Pratt, Rivest和 Tarjan在1973年开发的。如果我解释的不满意，他们1973年的论文肯定可以满足你。我没有使用文字介绍这个算法，而是使用配有很多注释的python代码来解释。

def pick_pivot(l):
    """
    Pick a good pivot within l, a list of numbers
    This algorithm runs in O(n) time.
    """
    assert len(l) > 0

    # If there are < 5 items, just return the median
    if len(l) < 5:
        # In this case, we fall back on the first median function we wrote.
        # Since we only run this on a list of 5 or fewer items, it doesn't
        # depend on the length of the input and can be considered constant
        # time.
        return nlogn_median(l)

    # First, we'll split `l` into groups of 5 items. O(n)
    chunks = chunked(l, 5)

    # For simplicity, we can drop any chunks that aren't full. O(n)
    full_chunks = [chunk for chunk in chunks if len(chunk) == 5]


    # Next, we sort each chunk. Each group is a fixed length, so each sort
    # takes constant time. Since we have n/5 chunks, this operation
    # is also O(n)
    sorted_groups = [sorted(chunk) for chunk in full_chunks]

    # The median of each chunk is at index 2
    medians = [chunk[2] for chunk in sorted_groups]

    # It's a bit circular, but I'm about to prove that finding
    # the median of a list can be done in provably O(n).
    # Finding the median of a list of length n/5 is a subproblem of size n/5
    # and this recursive call will be accounted for in our analysis.
    # We pass pick_pivot, our current function, as the pivot builder to
    # quickselect. O(n)
    median_of_medians = quickselect_median(medians, pick_pivot)
    return median_of_medians

def chunked(l, chunk_size):
    """Split list `l` it to chunks of `chunk_size` elements."""
    return [l[i:i + chunk_size] for i in range(0, len(l), chunk_size)]

让我们证明为什么 median-of-medians 是一个好的pivot。

红色圈起来的是chunks的中位数，中间的圆圈起来的是median-of-medians。回想一下，我们想让pivot尽可能地均分数组。考虑最坏的情况：pivot元素在数组的开头。

这张图可以分为四个象限：

• 左上：这个象限中的元素都比中位数小
• 左下：这个象限中的元素比中位数可能大，可能小
• 右上：这个象限中的元素也比中位数可能大，可能小
• 右下：这个象限中的元素都比中位数大

在这四个象限中左上和右下是有用的，因为它们中的元素要么大于要么小于median-of-medians，左下和右上没有用。

现在回到起初的问题上，最坏的情况就是pivot出现在列表的开头。正如我上面所说，至少左上象限的每个元素都小于pivot。总共有n个元素，每一列有5个元素，我们可以拿3个，同时可以拿一半的列，这样就有：

$$f(n)=\frac{3}{5}\times \frac{1}{2}\times n=\frac{3}{10}n$$

因此，每一步，我们至少可以移除30%的元素。

但是,每一步减少 30% 的元素够吗?在随机算法中平均要移除 50% 的元素。在每一步,我们的算法必须做到:

• $O(n)$时间来划分元素
• 使用当前数组的1/5来计算median of medians
• 当前数组的7/10来进行下一步的递归

总的运行时间如下所示：

$$T(n)=n+T(\frac{n}{5})+T(\frac{7n}{10})$$

要证明为什么是$O(n)$并不简单。我所知道的唯一方法是归纳法。

总结

有一个快速选择算法，当有一个足够好的pivot时，它可以在线性时间内找到数组的中位数。 还有一个median-of-medians算法，在时间复杂度为$O(n)$的情况下找到一个pivot（这个pivot对快速选择算法来说足够好）。将这两者结合，就有一个在线性时间内找到中位数（或者数组中的第K个元素）的算法。

实际中的线性时间查找中位数

在实际应用中，随机选择一个privot已经足够了。尽管median-of-medians方法也是线性的，在实际中它需要计算才可以得到。

最后，下面是每个实现要考虑到的元素数量的比较。这不是运行时的性能，而是quickselect函数所查看的元素总数(这里的元素总数应该就是递归次数)。它不包含计算median-of-medians。这张图的重点不是证明median-of-medians是一个好的算法，而是为了证明它是一个选择pivot的有效方法。