Python中二维列表中的二进制搜索和返回列表

Question

我是Python编程的新手，很长一段时间以来我一直在尝试这个问题，但是我无法提出解决方案。问题是，给了我一个雇员ID列表和雇员的出生年份，并且我应该返回一个在输入年份中出生的雇员ID列表。以下是列表的示例。

lst = [[2，1987]，[4，1990]，[6，1992] ...]

线性搜索需要太多时间，因为列表中有100万员工。因此，我认为这个问题要求我们使用二进制搜索。我应该想出一个函数，该函数需要输入年份和lst来得出返回员工ID的列表。如果该年没有雇员出生，则应返回一个空列表。以下是该功能的示例。

input> get_IDs（2012）

输出> [102、204、199632、199649]

我知道我可以对1D数组执行内置的bisect函数，但是我不确定如何对2D数组进行二进制搜索。请就我该怎么做提出建议，任何建议都将不胜感激！

Answer 1

首先，在现代硬件上一百万个元素并不多，因此使用过滤/列表可能会很快。让我们先设置一些测试数据：

import random, bisect, operator
random.seed(0)

years = list(range(1950, 2003))
N = 1_000_000

employees = sorted([(i, random.choice(years)) for i in range(N)],
                  key=operator.itemgetter(1))

target_year = 1980
%timeit emp_1980 = [i for i, y in employees if y == target_year]
# 1 loop, best of 3: 261 ms per loop
# can query 4 years per second

您可以将内置bisect与列表一起使用，而不是与标量一起使用，但是默认情况下，它将比较第一个元素（ID），而不是我们想要的年份。我们可以通过一些预处理来解决这个问题：

import bisect
# all (year, id) pairs sorted by year
year_id_sorted = [[y, i] for i, y in sorted(employees, key=operator.itemgetter(1))]

def ids_from_year(target_y):
  result = list()
  # make use of elementwise list ordering
  i = bisect.bisect_left(year_id_sorted, [target_y, 0])
  for y, ID in year_id_sorted[i:]:
    if y == target_y:
      result.append(ID)
    else:
      break
  return result

%timeit ids_from_year(target_year)
# 100 loops, best of 3: 11.9 ms per loop
# can query 100 years per second

这将使速度提高26倍。还不错，但是我们已经承担了一些预处理费用，并且必须存储数据集的副本。现在，让我们尝试用以下形式的字典存储员工：“年份=>员工列表”。

from collections import defaultdict
employee_by_year = defaultdict(list)
for i, y in employees:
  employee_by_year[y].append(i)
# 1 loop, best of 3: 361 ms per loop
# pre-processing step is only %40 more expensive than
# a single lookup in the original case.

%timeit employee_by_year[target_year]
# 10000000 loops, best of 3: 63.2 ns per loop
# can query 16 million years per second
# other factors will likely limit processing speed