二维数组Python的二分搜索算法

Question

我已经学会了进行二分搜索的一维数组方式：

def exist(target, array):
  lower = -1 
  upper = len(array)
  while not (lower + 1 == upper):
    mid = (lower + upper)//2
    if target== array[mid]: 
      return True
    elif target< array[mid]: 
      upper = mid 
    else: 
      lower = mid 
  return False

但是，现在我面临这个包含employee_id 和employee_birthyear 的二维数组list2：

[[2, 1986],
 [4, 1950],
 [6, 1994],
 [9, 2004],
 [12, 1988],
 [13, 1964],
 [16, 1987],
 [18, 1989],
 [19, 1951],
 [20, 1991]]

我想使用上面的纯二进制搜索算法编写一个函数，该算法以一年（作为整数）和list2 为参数，并返回具有匹配employee_birthyear 的employee_id 列表。

我该怎么做？

这是我想到的：

lst2 = [ j for i in list2 for j in i]

def get_IDs_with_birthyear(year, lst2):
    lower = -1 
  upper = len(lst2)
  while not (lower + 1 == upper): 
    mid = (lower + upper)//2
    if year = lst2[mid]:
      return mid 
  return []

更新： 我尝试对我的year 进行排序并进行二进制搜索，但是当同一年份有多个id 时，我无法检索所有ID。

ids = [] 
def get_IDs_with_birthyear(year, employee_with_birthyear_list):
  data2 = sorted(employee_with_birthyear_list, key=lambda d: d[1])
  years = [d[1] for d in data2]
  id = [d[0] for d in data2]
  
  lower = -1 
  upper = len(years)
  while not (lower + 1 == upper): 
    mid = (lower + upper)//2
    if year == years[mid]:
      ids.append(id[mid])
      return ids
    elif year < years[mid]:
      upper = mid 
    else: 
      lower = mid
  return False 

结果应该得到 [101, 201, 1999632, 1999649]，但我只得到 [1999632]

result = get_IDs_with_birthyear(1949, ewby)

我在我的函数中做错了什么？

Answer 1

二分搜索算法需要根据您正在搜索的键对数据进行排序。在这种情况下，您的 2D 列表不是按年份排序的，因此您需要采用不同的方法（或者您需要根据年份而不是员工 ID 对列表进行排序）。

一种选择是构建一个以年份为键的字典：

years = dict()
for emp,birth in employee_birthyear:
    years.setdefault(birth,[]).append(emp)

然后您可以获得任何给定出生年份的员工 ID 列表：

years[1950] # [4]

请注意，您会获得一份员工 ID 列表，因为您可能有多个同一年出生的员工

一旦构建字典（O(N) 操作而不是 O(NlogN) 排序），所有按年份的访问都将在 O(1) 中返回。这避免了更改employee_birthyear 列表中元素顺序的需要，并且比二分查找复杂度更低（O(logN)）

Answer 2

您可以利用bisect module 进行二进制搜索，并维护排序列表。

要按日期搜索，需要按日期对数据进行排序，并构建第二个键（日期）列表。然后，bisect_left() 可用于查找记录的开始 >= 选择的开始日期。

>>> from bisect import bisect_left, bisect_right
>>>
>>> data  = [[2, 1986],[4, 1950],[6, 1994],[9, 2004],[12, 1988],
...          [13, 1964],[16, 1987],[18, 1989],[19, 1951],[20, 1991]]
>>>
>>> # Sort data by date, then construct array of just dates.
>>> data2 = sorted(data, key=lambda d: d[1])
>>> dates = [d[1] for d in data2]
>>>
>>> # Use bisect_left() to locate the start of records >= 1988.
>>> idx = bisect_left(dates, 1988)
>>> data2[idx]     # Use that index with data2 to get first record.
[12, 1988]
>>>
>>> # Iteratively get records between 1988 and 1991.
>>> idx = bisect_left([d[1] for d in data2], 1988)
>>> recs = []
>>> while idx < len(data2) and data2[idx][1] <= 1991:
...     recs.append(data2[idx])
...     idx += 1
...     
>>> recs
[[12, 1988], [18, 1989], [20, 1991]]
>>> 
>>> 

这是一个灵活的解决方案，因为您要查找的日期实际上可能不在任何记录中，并且字典查找将失败。通过获取第一条记录 >= 日期，您可以向前迭代以按升序查找后续记录。您还可以使用 bisect_right() 指定结束日期并在数组中找到它的位置，这样您就可以拥有一系列感兴趣的记录。

>>> # A function that returns records that fall between 
>>> # specific dates, inclusive.
>>> def get_range(sorted_data, year_begin, year_end):
...     dates     = [d[1] for d in sorted_data]
...     idx_start = bisect_left(dates, year_begin)
...     idx_end   = bisect_right(dates, year_end)
...     return sorted_data[idx_start:idx_end]
...     
>>> get_range(data2, 1988, 1994)
[[12, 1988], [18, 1989], [20, 1991], [6, 1994]]
>>>
>>> # Add some more records for 1988 & sort.
>>> data2.extend([[1, 1988], [3, 1988]])
>>> data2.sort(); data2.sort(key=lambda r: r[1])
>>> 
>>> # Retrieve all records with 1988.
>>> get_range(data2, 1988, 1988)
[[1, 1988], [3, 1988], [12, 1988]]

如果数据集相对较小，而您只想获取具有特定日期的记录，则不需要排序或使用上述任何方法。您可以只过滤数据集。

>>> list(filter(lambda d: d[1] == 1988, data2))
[[1, 1988], [3, 1988], [12, 1988]]