Фильтрация координат по расстоянию от точки

У меня есть два массива:

A = np.array([[ 1.  ,  1.  ,  0.5 ],
              [ 2.  ,  2.  ,  0.7 ],
              [ 3.  ,  4.  ,  1.2 ],
              [ 4.  ,  3.  ,  2.33],
              [ 1.  ,  2.  ,  0.5 ],
              [ 6.  ,  5.  ,  0.3 ],
              [ 4.  ,  5.  ,  1.2 ],
              [ 5.  ,  5.  ,  1.5 ]])

B = np.array([2,1])

Я хотел бы найти все значения A, которые не находятся в радиусе 2 от B.

Мой ответ должен быть:

C = [[3,4,1.2],[4,3,2.33],[6,5,0.3],[4,5,1.2],[5,5,1.5]]

Есть ли питонический способ сделать это?

Что я пробовал:

radius = 2
C.append(np.extract((cdist(A[:, :2], B[np.newaxis]) > radius), A))

Но я понял, что np.extract сглаживает A, и я не получаю того, что ожидал.


person user3185171    schedule 11.12.2017    source источник


Ответы (2)


Пусть R будет радиусом здесь. У нас было бы несколько методов для ее решения, как обсуждается далее.

Подход 1. Использование cdist -

from scipy.spatial.distance import cdist

A[(cdist(A[:,:2],B[None]) > R).ravel()]

Подход 2. Использование np.einsum -

d = A[:,:2] - B
out = A[np.einsum('ij,ij->i', d,d) > R**2]

Подход №3: использование np.linalg.norm -

A[np.linalg.norm(A[:,:2] - B, axis=1) > R]

Подход № 4. Использование matrix-multiplication с np.dot -

A[(A[:,:2]**2).sum(1) + (B**2).sum() - 2*A[:,:2].dot(B) > R**2]

Подход № 5. Использование комбинации einsum и matrix-multiplication -

A[np.einsum('ij,ij->i',A[:,:2],A[:,:2]) + B.dot(B) - 2*A[:,:2].dot(B) > R**2]

Подход № 6: использование broadcasting -

A[((A[:,:2] - B)**2).sum(1) > R**2]

Следовательно, чтобы получить точки в радиусе R, просто замените > на < в вышеупомянутых решениях.

person Divakar    schedule 11.12.2017

Другой полезный подход, не упомянутый @Divakar, заключается в использовании cKDTree:

from scipy.spatial import cKDTree

# Find indices of points within radius
radius = 2
indices = cKDTree(A[:, :2]).query_ball_point(B, radius)

# Construct a mask over these points
mask = np.zeros(len(A), dtype=bool)
mask[indices] = True

# Extract values not among the nearest neighbors
A[~mask]

Основное преимущество заключается в том, что это будет намного быстрее, чем любой прямой подход, поскольку размер массива увеличивается, потому что структура данных позволяет избежать вычисления расстояния для каждой точки в A.

person jakevdp    schedule 12.12.2017