1 行中的双 for 循环以创建新的元组列表

Question

I have two sets of coordinates as numpy array. I would to create a new coordinate based on the first element of each set.

a = np.array([[1,2],[3,4],[5,6]])
b = np.array([[10,20],[30,40],[50,60]])

so I would like to get

[(1,10), (3,30), (5, 50)]

I tried:

c = [(i[0], j[0]) for (i,j) in itertools.product(a,b)]

But itreturned:

[(1, 10), (1, 30), (1, 50), (3, 10), (3, 30), (3, 50), (5, 10), (5, 30), (5, 50)]

which is a nested loop. Is there a one line solution here?

Many thanks

J

Answer 1

If you want result as tuples, first concatenate them along the second axis with np.c_ , view the result as an array of tuples and flatten the result with ravel :

np.c_[a[:,0],b[:,0]].view('i,i').ravel()
# array([(1, 10), (3, 30), (5, 50)], dtype=[('f0', '<i4'), ('f1', '<i4')])

---

Note that obtaining the result using numpy's vectorized operations will lead to huge benefits in terms of performance:

a = np.array([[1,2],[3,4],[5,6]])
b = np.array([[10,20],[30,40],[50,60]])
a = np.concatenate([a]*10000)
b = np.concatenate([b]*10000)

%timeit [(a_s[0], b_s[0]) for a_s, b_s in zip(a,b)]
# 19.2 ms ± 1.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit np.c_[a[:,0],b[:,0]].view('i,i').ravel()
# 82.1 µs ± 7.73 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Answer 2

Try the built-in function zip :

print([(a_s[0], b_s[0]) for a_s, b_s in zip(a,b)])
# [(1, 10), (3, 30), (5, 50)]

Answer 3

temp_arr = np.concatenate([a[:,:,None],b[:,:,None]],axis=2)
temp_arr[:,0]

array([[ 1, 10],
       [ 3, 30],
       [ 5, 50]])