Numpy: 提取子数组的高效方法

Question

我有一个很大的二维numpy数组 A，它的大小大约是 250000 x 30，还有两个一维的numpy数组 x 和 y。我想从 A 中提取出一个子数组，这个子数组的行是 x 中的行，列是 y 中的列。请问哪种方法更高效呢？

1. A[x[:,np.newaxis], y]
2. A[np.ix_(x,y)]

Answer 1

来做个基准测试吧！

import numpy as np

# some data
A = np.random.random((250000, 30))

# some random indices
x = np.random.randint(0, 250000, 150000)
y = np.random.randint(0, 30, 10)

def method1(A, x, y):
    return A[x[:, np.newaxis], y]

def method2(A, x, y):
    return A[np.ix_(x,y)]

def method3(A, x, y):
    return A[x][:,y]

def method4(A, x, y):
    return A[:,y][x]

这三种方法的基准测试结果如下：

method1: 87.7 ms
method2: 89.2 ms
method3: 115 ms
method4: 141 ms

所以，答案是这两种方法之间其实没有什么实质性的区别。