In [2]: def test_0(a, b):
   ...:     return a.shape[0] - np.searchsorted(a[::-1], b, side='right')
   ...:

In [3]: def test_1(a, b, opt):
   ...:     sort_idx = np.argsort(b)
   ...:     idx = np.zeros(b.shape[0], dtype=np.int64)
   ...:     if opt == 0:
   ...:         loc = np.searchsorted(b[sort_idx], a[::-1], side='left')
   ...:         k = a.shape[0] - 1
   ...:         for s, e in zip(loc[:-1], loc[1:]):
   ...:             idx[sort_idx[s:e]] = k
   ...:             k -= 1
   ...:         idx[sort_idx[:loc[0]]] = a.shape[0]
   ...:         idx[sort_idx[loc[-1]:]] = 0
   ...:     else:
   ...:         idx[sort_idx] = ((np.arange(b.shape[0]).reshape(-1, 1) - np.searchsorted(b[sort_idx], a[::-1], side='left')) < 0).sum(axis=1)
   ...:     return idx
   ...:

In [4]: @jit
   ...: def test_2(a, b):
   ...:     idx = np.zeros(b.shape[0], dtype=np.int64)
   ...:     for i in range(b.shape[0]):
   ...:         if b[i] >= a[0]:
   ...:             idx[i] = 0
   ...:         elif b[i] >= a[1]:
   ...:             idx[i] = 1
   ...:         elif b[i] >= a[2]:
   ...:             idx[i] = 2
   ...:         elif b[i] >= a[3]:
   ...:             idx[i] = 3
   ...:         elif b[i] >= a[4]:
   ...:             idx[i] = 4
   ...:         elif b[i] >= a[5]:
   ...:             idx[i] = 5
   ...:         elif b[i] >= a[6]:
   ...:             idx[i] = 6
   ...:         elif b[i] >= a[7]:
   ...:             idx[i] = 7
   ...:         elif b[i] >= a[8]:
   ...:             idx[i] = 8
   ...:         elif b[i] >= a[9]:
   ...:             idx[i] = 9
   ...:         elif b[i] >= a[10]:
   ...:             idx[i] = 10
   ...:         elif b[i] >= a[11]:
   ...:             idx[i] = 11
   ...:         elif b[i] >= a[12]:
   ...:             idx[i] = 12
   ...:         elif b[i] >= a[13]:
   ...:             idx[i] = 13
   ...:         elif b[i] >= a[14]:
   ...:             idx[i] = 14
   ...:         elif b[i] >= a[15]:
   ...:             idx[i] = 15
   ...:         elif b[i] >= a[16]:
   ...:             idx[i] = 16
   ...:         elif b[i] >= a[17]:
   ...:             idx[i] = 17
   ...:         elif b[i] >= a[18]:
   ...:             idx[i] = 18
   ...:         elif b[i] >= a[19]:
   ...:             idx[i] = 19
   ...:         else:
   ...:             idx[i] = 20
   ...:     return idx




In [6]: a = np.sort(np.linspace(200, 1000, 20))[::-1]

In [7]: b = np.random.randn(100000) * 1000

In [8]: np.allclose(test_0(a, b), test_1(a, b, 0))
Out[8]: True

In [9]: np.allclose(test_0(a, b), test_1(a, b, 1))
Out[9]: True

In [10]: np.allclose(test_0(a, b), test_2(a, b))
Out[10]: True




In [15]: %timeit -n 1000 -r 7 test_0(a, b)
1.49 ms ± 21.1 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [16]: %timeit -n 1000 -r 7 test_1(a, b, 0)
7.92 ms ± 47.8 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [17]: %timeit -n 1000 -r 7 test_1(a, b, 1)
14.3 ms ± 63.8 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

In [18]: %timeit -n 1000 -r 7 test_2(a, b)
1 ms ± 13.7 μs per loop (mean ± std. dev. of 7 runs, 1000 loops each)



【 在 CrTn 的大作中提到: 】
: np.searchsorted
--
FROM 68.197.87.*