Нахождение прямоугольников, охватывающих точки ⇐ Python

[Anonymous]

У меня есть список из n 2D-точек (x, y) и список из m прямоугольников (xmin, ymin, xmax, ymax). Я хочу найти индексы всех прямоугольников, окружающих точку. Это должно быть сделано эффективно. Также доступен графический процессор. Я рассматриваю решение, которое сильно распараллелено с помощью векторизованных операций. Для меня циклы for не подходят, так как 0 < n,m < 10^6
Например:

Код: Выделить всё

points = [[1, 1], [2, 2], [5, 5]]
rectangles = [[0, 0, 3, 3], [1, 1, 4, 4], [4, 4, 7, 7]]

result:
(1, 1): [(0, 0, 3, 3), (1, 1, 4, 4)]
(2, 2): [(0, 0, 3, 3), (1, 1, 4, 4)]
(5, 5): [(4, 4, 7, 7)]

Я написал это решение в PyTorch (попробовал разбить на фрагменты, чтобы уменьшить переполнение памяти)

Код: Выделить всё

def find_rectangles_containing_points(points, rectangles, chunk_size=1000):
n = points.size(0)
m = rectangles.size(0)

rectangles_containing_points = []

for i in range(0, n, chunk_size):
points_chunk = points[i:i + chunk_size]
points_expanded = points_chunk.unsqueeze(1)  # chunk_size x 1 x 2
rectangles_expanded = rectangles.unsqueeze(0)  # 1 x m x 4

is_inside = (rectangles_expanded[:, :, :2] 

Подробнее здесь: [url]https://stackoverflow.com/questions/77697515/finding-rectangles-enclosing-points[/url]