使用NumPy從另一個數組及其索引創建一個2D數組

Question

給定一個數組：

arr = np.array([[1, 3, 7], [4, 9, 8]]); arr

array([[1, 3, 7],
       [4, 9, 8]])

鑒於其指數：

np.indices(arr.shape)

array([[[0, 0, 0],
        [1, 1, 1]],

       [[0, 1, 2],
        [0, 1, 2]]])

我怎樣才能將它們整齊地疊在一起形成一個新的2D陣列？ 這就是我想要的：

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

這是我目前的解決方案：

def foo(arr):
    return np.hstack((np.indices(arr.shape).reshape(2, arr.size).T, arr.reshape(-1, 1)))

它有效，但有什么更短/更優雅的方式來執行此操作？

Answer 1

使用array-initialization然后broadcasted-assignment在后續步驟中分配索引和數組值 -

def indices_merged_arr(arr):
    m,n = arr.shape
    I,J = np.ogrid[:m,:n]
    out = np.empty((m,n,3), dtype=arr.dtype)
    out[...,0] = I
    out[...,1] = J
    out[...,2] = arr
    out.shape = (-1,3)
    return out

請注意，我們避免使用np.indices(arr.shape) ，這可能會減慢速度。

樣品運行 -

In [10]: arr = np.array([[1, 3, 7], [4, 9, 8]])

In [11]: indices_merged_arr(arr)
Out[11]: 
array([[0, 0, 1],
       [0, 1, 3],
       [0, 2, 7],
       [1, 0, 4],
       [1, 1, 9],
       [1, 2, 8]])

---

性能

arr = np.random.randn(100000, 2)

%timeit df = pd.DataFrame(np.hstack((np.indices(arr.shape).reshape(2, arr.size).T,\
                                arr.reshape(-1, 1))), columns=['x', 'y', 'value'])
100 loops, best of 3: 4.97 ms per loop

%timeit pd.DataFrame(indices_merged_arr_divakar(arr), columns=['x', 'y', 'value'])
100 loops, best of 3: 3.82 ms per loop

%timeit pd.DataFrame(indices_merged_arr_eric(arr), columns=['x', 'y', 'value'], dtype=np.float32)
100 loops, best of 3: 5.59 ms per loop

注意：計時包括轉換為pandas數據幀，這是此解決方案的最終用例。

Answer 2

nd數組的更通用的答案，正確處理其他dtypes：

def indices_merged_arr(arr):
    out = np.empty(arr.shape, dtype=[
        ('index', np.intp, arr.ndim),
        ('value', arr.dtype)
    ])
    out['value'] = arr
    for i, l in enumerate(arr.shape):
        shape = (1,)*i + (-1,) + (1,)*(arr.ndim-1-i)
        out['index'][..., i] = np.arange(l).reshape(shape)
    return out.ravel()

這將返回一個帶有索引列和值列的結構化數組，該列可以是不同類型的。