Python面试题系列之18: 查找最大或最小的N个元素？

Question

怎样从一个集合中获得最大或者最小的 N 个元素列表？

知识点详解

我们来分析下这道题。从一个集合中获得最大或者最小的 N 个元素列表，肯定就需要先对集合进行一个排序。

而要获取N个元素，N是不确定的，我们要分情况来讨论。因为不同场景的下最优解可能是不同的。

如下给定的集合X：

x = [3, 4, 2, 6, 8, 8, 9, 5, 3, -4, 7, -9, 11, 5, 32, 6, -52, 7, 9, 10]

当N = 1时

这时，题目就简化为从一个集合中获得最大值 或 最小值 ，那么我们可以直接使用内置函数max()、min()来解决。

>>> max(x)
32
>>> min(x)
-52

若想获取多个元素的话，max()、min()这俩函数就不适用了。

当N不等于1时

分析题目后，我们知道首先得对集合进行排序，然后再获取元素。提到排序，肯定就想到的sort()、sorted()这两个内置函数，（关于这俩函数的区别，之前题目已讲过，不再重复），我们这里直接选用sorted()来搞定排序这一步。

接下来就是获取一段连续元素，那么我们可以使用切片操作。

当N = 5时，集合中，最大的5个元素、最小的5个元素分别如下：

>>> print(sorted(x)[-5:])
[32, 11, 10, 9, 9]

>>> print(sorted(x)[:5])
[-52, -9, -4, 2, 3]

我们换一个思路。

在上一篇面试题中，我们用堆实现了一个优先级队列。而堆数据结构本身就是一个有序的，因为在底层实现里面，首先会先将集合数据进行堆排序后放入一个列表中，其中一个重要的特征是heap[0]永远是最小的元素。

在heap模块中给我们提供的nlargest()、nsmallest()这俩方法正好就能解决这个问题。

>>> heapq.__all__
['heappush', 'heappop', 'heapify', 'heapreplace', 'merge', 'nlargest', 'nsmallest', 'heappushpop']

同上当N = 5时，集合中最大的5个元素、最小的5个元素分别如下：

>>> print(heapq.nlargest(5, x))
[32, 11, 10, 9, 9]

>>> print(heapq.nsmallest(5, x))
[-52, -9, -4, 2, 3]

由上可见，两种思路都解决了这个问题，现在我们以取最大的N个元素为例，看看它们的效率如何。

效率

当N远小于集合的长度时，如：N = 2

>>> %timeit -n 10000 100 in heapq.nlargest(2, x)
5.26 µs ± 247 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit -n 10000 100 in sorted(x)[-2:]
910 ns ± 66.4 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

当N接近集合的一半长度时，如：N = 9

>>> %timeit -n 10000 100 in heapq.nlargest(9, x)
9.31 µs ± 272 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit -n 10000 100 in sorted(x)[-9:]
970 ns ± 47.5 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

当N接近集合的长度时，如：N = 19

>>> %timeit -n 10000 100 in heapq.nlargest(19, x)
9.03 µs ± 207 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

>>> %timeit -n 10000 100 in sorted(x)[-19:]
1.22 µs ± 77.3 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

综上对比可见，使用 sorted(items)[:N] / sorted(items)[-N:] 的方法要更优于 nlargest() / nsmallest() 方法。

扩展

此外，这两组函数都可以接受一个关键字参数，以获取用这个字段进行排序后的N个元素。

portfolio = [
    {'name': 'IBM', 'shares': 100, 'price': 91.1},
    {'name': 'AAPL', 'shares': 50, 'price': 543.22},
    {'name': 'FB', 'shares': 200, 'price': 21.09},
    {'name': 'HPQ', 'shares': 35, 'price': 31.75},
    {'name': 'YHOO', 'shares': 45, 'price': 16.35},
    {'name': 'ACME', 'shares': 75, 'price': 115.65}
]

使用 nlargest() / nsmallest() 方法来实现：

cheap = heapq.nsmallest(3, portfolio, key=lambda s: s['price'])
expensive = heapq.nlargest(3, portfolio, key=lambda s: s['price'])

使用sorted()方法来实现:

cheap = sorted(portfolio, key=lambda x: x['price'])[:3]
expensive = sorted(portfolio, key=lambda x: x['price'])[-3:]

Answer

当N=1时，
即查找最小或最大元素，这时直接使用内置函数min()和max()就好。

当N不等于1时，
有两种解决思路：sorted(items)[:N]/sorted(items)[-N:]和nlargest()/nsmallest()

从实际效率来看，前者更优。

后记

本篇中我们通过两种不同的思路解决一个集合中获得最大或者最小的 N 个元素列表的问题，你还有什么好的方法吗？欢迎在评论区留言分享~
好了，以上就是本篇全部内容。

备注：本篇首发于知识星球「人人都是Pythonista」。