首页 > 解决方案 > Python中二维列表的二进制搜索和返回列表

问题描述

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

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

线性搜索需要太多时间,因为列表中有 100 万员工。因此,我认为这个问题需要我们使用二进制搜索。我应该想出一个函数,该函数接受输入年份,并想出一个返回员工 ID 的列表。如果该年没有员工出生,它应该返回一个空列表。以下是该函数的示例。

输入 >get_IDs(2012)

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

我知道我可以为一维数组执行内置的 bisect 函数,但我不确定如何对二维数组进行二进制搜索。请就我应该做什么提出建议,任何建议都非常感谢!

标签: pythonpython-3.xlistmultidimensional-arraybinary-search

解决方案


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

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

我想最佳实现取决于您计划运行查询的频率。运行它至少两次证明使用dict. 如果您使用的是不那么“粒度”的指标(例如薪水),则可以使用二分搜索来证明其合理性。如果您需要查询多个指标,例如年份薪水,您会遇到内存与速度的权衡。那么,最佳策略实际上取决于您的数据分布。


推荐阅读