numpy二分查找

    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 Counter

import numpy as np

a = []
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 np

def 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 x

a = np.array([1, 2, 2, 4])
print(bisearch(a, 2, True, 'left'))#1
print(bisearch(a, 2, True, 'right'))#2
print(bisearch(a, 2, False, 'left'))#0
print(bisearch(a, 2, False, 'right'))#3
print(bisearch(a, -1, True, 'left'))#-1
print(bisearch(a, 5, True, 'right'))#4