Mapping every element of a 2D array to its corresponding (x,y) coordinate in python

user_2495203 0
Viewed 285

I have two 1D arrays that give a range of values for x and y -- for example:

x = np.array([0,1,2])
y = np.array([8,9])

I also have a corresponding 2D array that gives a z value associated with every possible (x,y) combination -- for example:

z = np.array([['A','B'],
              ['C','D'],
              ['E','F']])

Note how in this example z is a 3x2 array with each row corresponding to a given value of x and each column corresponding to a given value of y (order is important for this mapping).

What I want is to combine those 3 arrays into a 1D list of coordinate-value pairs -- for example:

result = [(0,8,'A'),
          (0,9,'B'),
          (1,8,'C'),
          (1,9,'D'),
          (2,8,'E'),
          (2,9,'F')]

There is definitely a brute force way of doing this using for loops, but is there an easier/faster way using pre-existing numpy or python routines? My guess is it might involve np.meshgrid(), zip() or itertools but I can't figure it out. Any help would be greatly appreciated!

itertools tuples indexing numpy python
3 Answers

Try using itertools.product with zip and list comprehension. Also you need to unstack your 2-d list using numpy

from itertools import product
import numpy as np

x = np.array([0,1,2])
y = np.array([8,9])
z = np.array([['A','B'],
              ['C','D'],
              ['E','F']])

result = [(x,y,z) for (x,y),z in zip(product(x, y), np.hstack(z))]
# [(0, 8, 'A'), (0, 9, 'B'), (1, 8, 'C'), (1, 9, 'D'), (2, 8, 'E'), (2, 9, 'F')]
user_91778770

You can use itertools.product and zip the product with flattened z:

from itertools import product
out = [tuple(list(i)+[j]) for i,j in zip(product(x,y),z.flatten())]

Output:

[(0, 8, 'A'), (0, 9, 'B'), (1, 8, 'C'), (1, 9, 'D'), (2, 8, 'E'), (2, 9, 'F')]
Admin1

You can use np.meshgrid as well, that's faster than itertools.product:

combs = np.stack(np.meshgrid(x, y), axis=-1).reshape(-1, 2)
print(combs)
[[0 8]
 [1 8]
 [2 8]
 [0 9]
 [1 9]
 [2 9]]

It's not a good idea to keep data of different type in the same array. So replace it with dynamic types in a cost of iteration if you really need it:

[(*c, id) for c, id in zip(combs, z.ravel())]
>>> [(0, 8, 'A'), (1, 8, 'B'), (2, 8, 'C'), (0, 9, 'D'), (1, 9, 'E'), (2, 9, 'F')]
user_30448250
Related