根据 x 个其他数组对 numpy 数组进行切片

Question

我正在尝试使我当前的代码动态化。这意味着无论用户的数组输入数量如何，它都应该能够进行调整。

当前代码：

main1 = numpy.array([1,2,3,4])
array1 = numpy.array(['a','b','c','b'])
my_list1 = ['a','b']
array2 = numpy.array(['cat','dog','bird','cat'])
my_list2 = ['cat']

result_array = main1[np.in1d(array1, my_list1) and np.in1d(array2, my_list2)]

打印出result_array的期望结果是：

array([1, 4])

这是因为a 和cat & b 和cat 的交集。

我的目标是能够使用 n 个 array1、array2 ... 和 n 个 my_list1、my_list2... 来做到这一点。

提前致谢！

Answer 1

两个以上数组的版本，使用logical_and.reduce：

array3 = numpy.array(['cat3','dog3','bird3','cat3'])
my_list3 = ['cat3']

my_arrays = [array1, array2, array3]
my_lists = [my_list1, my_list2, my_list3]
res1 = main1[numpy.logical_and.reduce(tuple(np.in1d(array, lst) for 
                                            array, lst in zip(my_arrays, my_lists)))]

测试一下：

res2 = main1[np.in1d(array1, my_list1) & np.in1d(array2, my_list2) &
             np.in1d(array3, my_list3)]

看起来不错：

>>> np.all(res1 == res2)
True

旧答案只有两个数组。

这应该可行：

my_arrays = [array1, array2]
my_lists = [my_list1, my_list2]
main1[np.logical_and(*(np.in1d(array, lst) for array, lst in zip(my_arrays, my_lists)))]

结果：

array([1, 4])