a = np.array([1, 2, 2, 3]) print(np.searchsorted(a, 0)) # 0 print(np.searchsorted(a, 1)) # 0 print(np.searchsorted(a, 2)) # 1 print(np.searchsorted(a, 2, 'left')) # 1 print(np.searchsorted(a, 2, 'right')) # 3 print(np.searchsorted(a, 2.5, 'right')) # 3 print(np.searchsorted(a, 2.5, 'left')) # 3 print(np.searchsorted(a, 3, 'left')) # 3 print(np.searchsorted(a, 3, 'right')) # 4 print(np.searchsorted(a, 4)) # 4 print(np.searchsorted(a, [0, 1, 2, 3, 4, 5, ])) # [0 0 1 3 4 4]
searchsorted有三个重要参数：
- a：待查找的有序数组
- v：待查找的值
- side：字符串，取值为left或者right，表示取下界（闭区间）还是取上界（开区间），默认参数为下界闭区间
利用searchsorted可以非常炫酷地实现轮盘赌随机选取：
t = np.cumsum(weights) sample = np.searchsorted(t, np.random.random() * t[-1])
cumsum保证了递增，searchsorted二分查找，其中t[-1]表示全部元素之和，整个过程一气呵成、美不胜收。
虽然如此，这种方式依然不是最好的方法。因为numpy提供了轮盘赌算法。from collections import Counterimport numpy as npa = []for i in range(10000): x = np.random.choice([1, 2, 3], 2, p=[0.1, 0.3, 0.6]) a.extend(x)a = Counter(a)a = np.array([np.array(i) for i in a.items()], dtype=np.float32)a[:, 1] /= np.sum(a[:, 1])print(a)
输出为
[[1. 0.0993 ] [2. 0.30325] [3. 0.59745]]
因为searchsorted的下闭区间、上开区间效果有些奇特，所以可以wrap一下使它的行为更加明确
二分查找实际上可以写成四种：
- 左闭区间
- 右闭区间
- 左开区间
- 右开区间
如果自己写，一定要十分小心地考虑好边界条件才能够避免出错。
import numpy as npdef bisearch(a, v, can_eq=True, side='left'): x = np.searchsorted(a, v, side=side) if x >= a.shape[0]: return x if can_eq: if side == 'left': if a[x] == v: return x else: return x - 1 else: if a[x] > v: if x > 0 and a[x - 1] == v: return x - 1 else: return x else: return x else: if side == 'left': if a[x] == v: return x - 1 else: return x else: return xa = np.array([1, 2, 2, 4])print(bisearch(a, 2, True, 'left'))#1print(bisearch(a, 2, True, 'right'))#2print(bisearch(a, 2, False, 'left'))#0print(bisearch(a, 2, False, 'right'))#3print(bisearch(a, -1, True, 'left'))#-1print(bisearch(a, 5, True, 'right'))#4