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

Question

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!

Answer 1

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')]