Organize 3d points and found them by distance from a position

Question

I'm working in 3D context. I've some objects in this space who are represented by x, y, z position.

# My objects names (in my real context it's pheromone "point")
A = 1
B = 2
C = 3
D = 4

# My actual way to stock their positions
pheromones_positions = {
    (25, 25, 60): [A, D],
    (10, 90, 30): [B],
    (5, 85, 8): [C]
}

My objective is to found what points (pheromones) are near (with distance) a given emplacement. I do this simply with:

def calc_distance(a, b):
    return sqrt((a[0]-b[0])**2+(a[1]-b[1])**2+(a[2]-b[2])**2)

def found_in_dict(search, points, distance):
    for point in points:
        if calc_distance(search, point) <= distance:
            return points[point]

founds = found_in_dict((20, 20, 55), pheromones_positions, 10)
# found [1, 4] (A and D)

But, with a lot of pheromones it's very slow (test them one by one ...). How can i organize these 3D positions to found more quickly "positions by distance from given position" ? Does exist algorithms or librarys (numpy ?) who can help me in this way ?

Answer 1

You should compute all (squared) distances at once. With NumPy you can simply subtract the target point of size 1x3 from the (nx3) array of all position coordinates and sum the squared coordinate differences to obtain a list with n elements:

squaredDistances = np.sum((np.array(pheromones_positions.keys()) - (20, 20, 55))**2, axis=1)
idx = np.where(squaredDistances < 10**2)[0]
print pheromones_positions.values()[idx]

Output:

[1, 4]

---

By the way: Since your return statement is within the for-loop over all points, it will stop iterating after finding a first point. So you might miss a second or third match.