View ndarray as list of indexes and list of values

user_12326226 0
Viewed 41

I have an ND-array with 2 or 6 dimensions. For example:

nd_array = [[90, 80],
            [70, 60], 
            [50, 40]]

and I want to view it as a list of indices like so:

nd_array_transformed = [[90, 0, 0],
                        [80, 0, 1],
                        [70, 1, 0],
                        [60, 1, 1],
                        [50, 2, 0], 
                        [40, 2, 1]]

or a similar representation.

I found this code for the 2D case, but I'm looking for a more general solution:

np.array([[nd_array[row][col], row, col] for row in range(nd_array.shape[0]) 
                                         for col in range(nd_array.shape[1])])

Is it possible to make this transformation in the general case, perhaps using *nd_array.shape?

numpy-ndarray indexing python
1 Answers

You can use itertools.product() with range() and map() to generate all the possible indices, then use the unpacking operator to add the indices to each inner list. The example used here is two dimensions, but the use of unpacking as opposed to a fixed number of for loops makes this approach extensible to an arbitrary number of dimensions:

import numpy as np
from itertools import product
nd_array = np.array([[90, 80],
            [70, 60],
            [50, 40]])

result = np.array([[nd_array[indices], *indices]
    for indices in product(*map(range, nd_array.shape))])
print(result)

The syntax used to construct what we're iterating over in the list comprehension is a bit dense, so here's a more detailed explanation:

.shape produces a tuple that gives the dimensions of the array (e.g. (3, 2)). The examples given will all assume this shape.

We want to be able to generate all possible indices given by the shape (e.g. (0, 0), (0, 1) ... (2, 1)). We use map() with range(), taking in each integer in the tuple and producing the corresponding range() initialized with that integer (i.e. (3, 2) is transformed into an iterable containing range(3) and range(2)). This iterable is then unpacked (using *) into itertools.product() -- the function call essentially becomes itertools.product(range(3), range(2)), generating a Cartesian product that gives us all the desired indices.

user_177698150
Related