python - 根据numpy中的条件形成子数组

Question

我有一组值，例如

import numpy as np 
dats = np.array([[r1,x1,y1],[r2,x2,y2],...,[rn,xn,yn]])

我需要找到一个子数组，它的元素使得坐标 xi,yi 位于另一个点 xp,yp 周围半径为 R 的圆内。这是一个解决方案:

def dats_within_radius(R,xp,yp): 
    temp = np.array([])
    for loop in range(dats.shape[0]):
        r_loop = dats[loop][0]
        x_loop = dats[loop][1]
        y_loop = dats[loop][2]
        if (x_loop-x)**2+(y_loop-y)**2 <= R**2:
            temp = np.append(temp, [r_loop,x_loop,y_loop])
    temp = temp.reshape((int(temp.size/3),3))
    return temp

一定有更好的方法，对吧？谁能给我推荐一下？

Answer 1

这是一个使用 broadcasting 的矢量化单行代码-

def dats_within_radius_vectorized(R,xp,yp): 
    return dats[((dats - [R,xp,yp])**2).sum(1) <= R**2]

为了性能，我建议使用 np.einsum , 像这样 -

def dats_within_radius_vectorized2(dats, R,xp,yp): 
    subs = dats - [R,xp,yp]
    return dats[np.einsum('ij,ij->i',subs, subs) <= R**2]

---

运行时测试-

In [82]: # Inputs
    ...: dats = np.random.randint(0,9,(10000,3))
    ...: R,xp,yp = 5,2,1
    ...: 

In [83]: %timeit dats_within_radius(dats,R,xp,yp) # Original soln
    ...: %timeit dats_within_radius_loop_compr(dats,R,xp,yp) # @BloodyD's soln
    ...: %timeit dats_within_radius_vectorized(dats,R,xp,yp)
    ...: %timeit dats_within_radius_vectorized2(dats,R,xp,yp)
    ...: 
10 loops, best of 3: 117 ms per loop
10 loops, best of 3: 38.8 ms per loop
1000 loops, best of 3: 517 µs per loop
1000 loops, best of 3: 410 µs per loop