如何在二维数组中使用 numpy.searchsorted 进行矢量化答案

【问题标题】：How to vectorize with numpy.searchsorted in a 2d-array如何在二维数组中使用 numpy.searchsorted 进行矢量化
【发布时间】：2019-10-21 14:11:04
【问题描述】：

我有一个二维数组 (a) 用于查找，还有一个数组 (v) 用于查找应插入元素的索引：

import numpy as np

# [EDIT] Add more records which contain NaNs
a = np.array(
[[0., 923.9943, 996.8978, 1063.9064, 1125.639, 1184.3985, 1259.9854, 1339.6107, 1503.4462, 2035.6527],
 [0., 1593.6196, 1885.2442, 2152.956, 2419.0038, 2843.517, 3551.225, 5423.009, 18930.8694, 70472.4002],
 [0., 1593.6196, 1885.2442, 2152.956, 2419.0038, 2843.517, 3551.225, 5423.009, 18930.8694, 70472.4002],
 [0., 1084.8388, 1132.6918, 1172.2278, 1215.7986, 1259.062, 1334.4778, 1430.738, 1650.4502, 3966.1578],
 [0., 1084.8388, 1132.6918, 1172.2278, 1215.7986, 1259.062, 1334.4778, 1430.738, 1650.4502, 3966.1578],
 [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
 [0., 923.9943, 996.8978, 1063.9064, 1125.639, 1184.3985, 1259.9854, 1339.6107, 1503.4462, 2035.6527],
 [0., 1593.6196, 1885.2442, 2152.956, 2419.0038, 2843.517, 3551.225, 5423.009, 18930.8694, 70472.4002],
 [0., 1593.6196, 1885.2442, 2152.956, 2419.0038, 2843.517, 3551.225, 5423.009, 18930.8694, 70472.4002],
 [0., 1084.8388, 1132.6918, 1172.2278, 1215.7986, 1259.062, 1334.4778, 1430.738, 1650.4502, 3966.1578],
 [0., 1084.8388, 1132.6918, 1172.2278, 1215.7986, 1259.062, 1334.4778, 1430.738, 1650.4502, 3966.1578]])

v = np.array([641.954, 56554.498, 168078.307, 1331.692, 2233.327, 1120.03, 641.954, 56554.498, 168078.307, 1331.692, 2233.327])

这是我想要得到的结果：

[1, 9, 10, 6, 9, 0, 1, 9, 10, 6, 9]

显然，使用 for 循环我可以像这样索引数组 a 和 v：

for i, _ in enumerate(a):
    print(np.searchsorted(a[i], v[i]))

有没有更高效的vectorized 方法？

【问题讨论】：

那么，它可以将整行作为所有 NaN 吗？是否也可以将少量元素而不是所有元素连续作为 NaN？
@Divakar 不可能连续出现几个 NaN。如果它有 NaN，则整行都是 NaN。
考虑使用 Jax 的 vmap

标签： python arrays performance numpy search

【解决方案1】：

灵感来自Vectorized searchsorted numpy 的基本思想，这里是2D 和1D 数组之间的一个-

def searchsorted2d(a,b):
    # Inputs : a is (m,n) 2D array and b is (m,) 1D array.
    # Finds np.searchsorted(a[i], b[i])) in a vectorized way by
    # scaling/offsetting both inputs and then using searchsorted

    # Get scaling offset and then scale inputs
    s = np.r_[0,(np.maximum(a.max(1)-a.min(1)+1,b)+1).cumsum()[:-1]]
    a_scaled = (a+s[:,None]).ravel()
    b_scaled = b+s

    # Use searchsorted on scaled ones and then subtract offsets
    return np.searchsorted(a_scaled,b_scaled)-np.arange(len(s))*a.shape[1]

给定样本的输出 -

In [101]: searchsorted2d(a,v)
Out[101]: array([ 1,  9, 10,  6,  9])

所有 NaN 行的大小写

要扩展使其适用于所有 NaN 行，我们需要更多步骤 -

valid_mask = ~np.isnan(a).any(1)
out = np.zeros(len(a), dtype=int)
out[valid_mask] = searchsorted2d(a[valid_mask],v[valid_mask])

【讨论】：

非常感谢。这是一个好主意，并且与示例配合得很好，但是对于包含 NaN 的真实数据，出现了一些问题。你能给我一个处理这个的想法吗？
@kokon 请发布示例数据、预期输出，并希望提供一个带有循环的有效解决方案。
由于此评论中的字符数量有限，我更新了上面原始帖子中的示例以包含 NaN 记录。目前使用您的解决方案，NaN 记录后的数据返回负数，这与预期的输出不同： array([ 1, 9, 10, 6, 9, 0, -10, -20, -30, -40, - 50])