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

我有一组值，例如

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 

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

请您参考如下方法:

这是一个使用 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