通过多重索引/分组解开

Question

我有一个numpy数组，还有对应的行和列索引：

mat = np.array([[0, 1, 2], 
                [3, 4, 5], 
                [6, 7, 8]])

row_idx = np.array([0, 0, 0, 
                    0, 0, 0,
                    1, 1, 1])

col_idx = np.array([0, 1, 2,
                    0, 1, 2,
                    0, 1, 2])

我想根据row_idx和col_idx指定的分组来展开这个矩阵：

result = np.array([0, 3, 1, 4, 2, 5, 6, 7, 8])

我可以用numpy_groupies.aggregate()来做到这一点，但返回的结果有时会很奇怪：

import numpy as np
import numpy_groupies as npg

mat = np.array([[0, 1, 2], 
                [3, 4, 5], 
                [6, 7, 8]])

row_idx = np.array([0, 0, 0, 
                    0, 0, 0,
                    1, 1, 1])

col_idx = np.array([0, 1, 2,
                    0, 1, 2,
                    0, 1, 2])


result = npg.aggregate(group_idx=np.vstack([row_idx, col_idx]), a=mat.reshape(-1), func='array')

Produces: 

[[array([0, 3]) array([1, 4]) array([2, 5])]
 [array([6]) array([7]) array([8])]]

我试过使用numpy_groupies包里的aggregate()，但它既慢又会返回一个混合了np.arrays和整数的数组，在某些情况下还可能包含整数而不是数组（这让后续处理变得困难且缓慢）。

为什么ravel不管用

ravel的问题在于，它不能处理那种矩阵中有些区域是按行分组，有些区域是按列分组的情况，这种分组是基于row_idx和col_idx的。也就是说，可能会同时存在行和列的分组。