使用另一个数组重新排序和切片 numpy 数组答案

【问题标题】：Reordering and slicing a numpy array using another array使用另一个数组重新排序和切片 numpy 数组
【发布时间】：2021-06-23 18:42:44
【问题描述】：

给定两个数组arr1 和arr2，我想从arr1 中提取第一个cap 元素，按照它在arr2 中出现的顺序。所以从下面的例子：

arr1 = np.array([3, 7, 1, 10, 2])
arr2 = np.array([0, 3, 1, 5, 2, 7, 6, 4])
cap = 3

从arr1，3,7,1,2 出现在arr2 排序为3,1,2,7。我只想要前 3 个，所以我想要的结果是 array([3, 1, 2])。

我可以得到我想要的结果如下：

ans = np.array([el for el in arr2 if el in arr1][:cap])

但是当我只想要前 3 次出现时，此解决方案会遍历整个 arr2。有没有更有效的解决方案？

【问题讨论】：

这能回答你的问题吗？ How to limit the size of a comprehension?
如果您非常关心性能，您可能需要重新访问我修改后的答案并查看 Tom 的调整图，以了解它与替代方案相比的性能。显然，修改后的 for 循环在大多数情况下都超过了您原来的列表理解以及 np.isin()。

标签： python numpy

【解决方案1】：

这是一个不使用任何迭代的尝试，并将成员资格检查切换到np.isin()：

>>> arr2[np.isin(arr2, arr1[np.isin(arr1, arr2)])][:cap]
array([3, 1, 2])

由于这可能比您的列表理解更难阅读，因此以下是分解的步骤：

# find elements of arr1 in arr2
subset_arr1 = arr1[np.isin(arr1, arr2)] # array([3, 7, 1, 2])

# find the elements of arr2 in subset1, this gets the desired order
subset_arr2 = arr2[np.isin(arr2, subset_arr1)] # array([3, 1, 2, 7])

# take the first cap elements
ans = subset_arr2[:cap] # array([3, 1, 2])

时间

我最初的测试表明我的np.isin() 方法是最快的。然而，Georgy 更新了他们的答案（for 循环）并修复了错误。然后，新的测试表明 Georgy 的答案对于所有测试数据大小都是最快的：

“np.isin”是我的答案，“listcomp”是您的原始解决方案，“forloop”是other existing answer。在上面，N的大小（x轴）是arr1的长度，arr2是那个长度的两倍，cap固定在5（我以为这似乎是合理的，但请参阅下面的调整）。

似乎通过 for 循环提前停止的能力非常有益！在我的回答中，代码必须检查 arr1 的 all 元素是否属于 arr2（上面的 subset_arr1），然后检查 arr2 的 all 元素为该子集的成员资格。 Georgy 的循环 OTOH 只需检查arr1 中arr2 的一些元素的成员资格。所以即使我的循环应该更快，也可以有更多！

我尝试摆弄的唯一参数是cap，它显示了更小的cap 大小对于for 循环更好（更小的上限-> 更少的迭代）。以上是cap == 5； cap == 100:

还有cap == 1000：

我认为其他参数也会影响事情（例如，is-in 操作返回 true AKA 数值范围或两个数组的相对大小的可能性有多大）。如果有人想测试，我已经使下面的perfplot 代码更加参数化 - 我厌倦了考虑复杂性?：

import numpy as np
import perfplot

cap = 5
minnumber = 1
maxnumber = 10
factor1 = 1 # how size of arr1 relates to n
factor2 = 2 # how size of arr2 relates to n
low_n = 3
high_n = 18

def make_data(n):
    arr1 = np.random.randint(minnumber, maxnumber, n*factor1)
    arr2 = np.random.randint(minnumber, maxnumber, n*factor2)
    return arr1, arr2

def listcomp(arr1, arr2):
    return np.array([el for el in arr2 if el in arr1][:cap])

def npisin(arr1, arr2):
    return arr2[np.isin(arr2, arr1[np.isin(arr1, arr2)])][:cap]

def forloop(arr1, arr2):
    capcopy = cap
    lst=[]
    for el in arr2:
        if el in arr1:
            lst.append(el)
            capcopy -=1
        if capcopy <1:
            break
    return np.array(lst)

perfplot.show(
    setup=lambda n : make_data(n),
    kernels = [
        lambda x: listcomp(*x),
        lambda x: npisin(*x),
        lambda x: forloop(*x),
        ],
    labels=['listcomp', 'np.isin', 'forloop'],
    n_range=[2 ** k for k in range(low_n, high_n)])

【讨论】：

好建议！我稍微更新了我的 for 循环，因为在旧版本中，在达到上限后迭代实际上并没有停止，但后来的匹配根本不包括在内。知道这会如何改变事情会很有趣。
@GeorgyKopshteyn 您的更改绝对是有保证的，而且在许多情况下似乎比我的更快！ +1。看看我上面的想法

【解决方案2】：

这是一个简单的迭代方法，在达到 cap 后停止：

lst=[]
for el in arr2:
    if el in arr1:
        lst.append(el)
        cap -=1
    if cap <1:
        break
ans = np.array(lst)

【讨论】：