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

时间:2019-07-19 08:04:48

标签: python numpy

我有两组坐标作为numpy数组。我将基于每个集合的第一个元素创建一个新坐标。

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

所以我想得到

  

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

我尝试过:

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

但返回:

  

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

这是一个嵌套循环。这里有一种解决方案吗?

非常感谢

J

3 个答案:

答案 0 :(得分:5)

如果要将结果作为元组,请首先使用np.c_将其沿第二个轴连接,将结果作为元组数组查看,并使用ravel将结果展平:

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

请注意,使用numpy的矢量化操作获取结果将带来巨大的性能优势:

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)

答案 1 :(得分:3)

尝试使用内置功能zip

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

答案 2 :(得分:0)

temp_arr = np.concatenate([a[:,:,None],b[:,:,None]],axis=2)
temp_arr[:,0]
array([[ 1, 10],
       [ 3, 30],
       [ 5, 50]])
相关问题