diff --git a/poly_match_v1.py b/poly_match_v1.py
index 675c88a..4293a46 100644
--- a/poly_match_v1.py
+++ b/poly_match_v1.py
@@ -1,4 +1,5 @@
from functools import cached_property
+from itertools import compress
from typing import List, Tuple
import numpy as np
from dataclasses import dataclass
@@ -56,11 +57,8 @@ def generate_example() -> Tuple[List[Polygon], List[np.array]]:
def find_close_polygons(
polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
- close_polygons = []
- for poly in polygon_subset:
- if np.linalg.norm(poly.center - point) < max_dist:
- close_polygons.append(poly)
-
+ dist = np.linalg.norm([poly.center for poly in polygon_subset] - point, axis=1)
+ close_polygons = list(compress(polygon_subset, dist < max_dist))
return close_polygons
10x faster than the original, without resorting to native code, and without substantial increase in complexity of code base.
