如何有效地获得由另一个矩阵索引的子矩阵?

时间:2016-10-05 06:23:55

标签: performance pandas numpy matrix

现在我有一个矩阵:test_projs和一个索引矩阵,GT_index

GT_index每行的每个元素表示test_projs中该行所需元素的索引。现在我想计算选区上每一行的总和。

一个简单的例子:

test_projs = pd.DataFrame({0:[1,2,3], 1:[4,5,6], 2:[7,8,9], 3:[1,3,5]}, index=[3,4,5])
GT_index = pd.DataFrame({0:[3,2,1], 1:[2,3,2], 2:[1,1,3]}, index=[3,4,5])
b = np.zeros(3)
i = 0
for (index1, row1),(index2,row2) in zip(test_projs.iterrows(),GT_index.iterrows()) :
    b[i]=row1[row2].values.sum()
    i+=1
b

以上例子的结果是:

enter image description here

但它很慢。你知道如何提高效率吗?

1 个答案:

答案 0 :(得分:2)

这是一种使用NumPy数组和fancy-indexing技术的方法 -

arr1 = test_projs.values
arr2 = GT_index.values
out = arr1[np.arange(arr2.shape[0]),arr2.T].sum(0)