多维索引和映射

Question

我看到numpy 有许多索引工具，但仍然无法让它们完成我需要的工作。

首先，假设有两个形状相同的一维数组A, I，一个可以用I的元素索引的一维数组B，以及一个三参数函数f。然后我需要的结果可以像starmap(f, zip(A, I, B[I])) 一样实现（starmap 和zip 来自纯python，而不是numpy）。到目前为止，一切都很好......

但实际上，所有数组都是二维，我也想得到二维结果，这相当于将与上面相同的函数应用于数组的每一行 -这就是我现在循环执行的操作。

除了循环之外，还有更好的方法吗？

UPD：

例如，一维数组：

A = np.random.randint(0, 10, size=(3,))
B = np.random.randint(0, 10, size=(5,))
I = np.random.randint(0, 5, size=(3,))

def f(a, i, b):
    return (a, i, b)

print A, I, B
print list(starmap(f, zip(A, I, B[I])))

对于二维：

A = np.random.randint(0, 10, size=(2, 3))
B = np.random.randint(0, 10, size=(2, 5))
I = np.random.randint(0, 5, size=(2, 3))

def f(a, i, b):
    return (a, i, b)

print A
print I
print B
print [list(starmap(f, zip(A_row, I_row, B_row[I_row])))
       for A_row, I_row, B_row in zip(A, I, B)]

Answer 1

如果它只是一个普通函数，那么您的开销将主要是函数开销。您可以使用hstack 或其他方式执行zip，但这对速度没有帮助。

如果您只是在寻找一种更简洁的方法，请尝试

x_indexes, _ = numpy.ogrid[:len(I), :0]
numpy.vectorize(f)(A, I, B[x_indexes, I])

通常，更高级别的视图可以让您矢量化整个事物，这会更快。如果结果很慢，请记住这一点。