Convert n-dimensional numpy array to 2-dimensional index array

Question

I would like to convert an n-dimensional numpy array such as this one:

[ [ a, b, c],
  [ d, e, f] ]

to a 2-dimensional array of axis_0_index , axis_1_index , cell_value .

[ [ 0, 0, a],
  [ 0, 1, b],
  [ 0, 2, c],
  [ 1, 0, d],
  [ 1, 1, e],
  [ 1, 2, f] ]

Is this possible to do easily in NumPy?

Answer 1

You can use (abuse?) np.where to get all the indices of the array, using an array of ones of the same shape as condition, then stack those indices with the (flattened) array and finally transpose.

>>> A = np.array([ [ 'a', 'b', 'c'], [ 'd', 'e', 'f'] ])
>>> ones = np.ones(A.shape)
>>> np.vstack(np.where(ones) + (A.ravel(),)).transpose()
array([['0', '0', 'a'],
       ['0', '1', 'b'],
       ['0', '2', 'c'],
       ['1', '0', 'd'],
       ['1', '1', 'e'],
       ['1', '2', 'f']], 
      dtype='|S1')

---

After some more searching, its probably cleaner to use np.indices :

>>> X, Y = np.indices(A.shape)
>>> np.vstack((X.ravel(), Y.ravel(), A.ravel())).T

or

>>> np.vstack((X, Y, A)).reshape(3,A.size).T

The result, in both cases, is the same as above.

---

I did some timing analysis using IPython's %timeit . Curiously, my first solution with where seems to be fastest, at least for this very small test array:

>>> %timeit f1() # using ones and np.where
10000 loops, best of 3: 72.3 us per loop
>>> %timeit f2() # using np.indices and ravel
10000 loops, best of 3: 125 us per loop
>>> %timeit f3() # using np.indices and reshape
10000 loops, best of 3: 110 us per loop
>>> %timeit g() # using meshgrid
10000 loops, best of 3: 134 us per loop

Answer 2

You can use np.meshgrid as shown in the sample run below -

In [19]: A
Out[19]: 
array([[19, 80, 63],
       [24, 54, 44]])

In [20]: m,n = A.shape

In [21]: R,C = np.meshgrid(np.arange(m),np.arange(n))

In [22]: np.column_stack((R.ravel('F'),C.ravel('F'),A.ravel()))
Out[22]: 
array([[ 0,  0, 19],
       [ 0,  1, 80],
       [ 0,  2, 63],
       [ 1,  0, 24],
       [ 1,  1, 54],
       [ 1,  2, 44]])