Find the k closest points in R3

closest_k_in_r3.py

#!/usr/bin/env python

import heapq
from collections import namedtuple

Point = namedtuple('x y z')
PointState = namedtuple('dist point')


class ClosestKPoints(object):
    """ ClosestKPoints(k=3, N=[
            Point(0, 1, 2),
            Point(5, 9, 3),
            ...
        ])
    """
    
    def __init__(self, k, N):
        self.k = k
        self.heap = []
        self.furthest_dist = -1
        for p in N:
            self.offer(p)
    
    def distance(self, p):
        # sqrt doesn't effect outcome
        return p.x**2 + p.y**2 + p.z**2
    
    def offer(self, p):
        # heap invariant is maintained internally by heap.sort()
        # smallest element is at heap[0], so we'll negate the distance to keep biggest element on top
        dist = -1 * self.distance(p)
        if len(self.heap) < self.k:
            self.heap.append(PointState(dist, p))
            self.furthest_dist = max(self.furthest_dist, dist)
        elif dist > self.furthest_dist:
            if len(self.heap) is self.k:
                heapq.heapify(self.heap)
            heapq.heappushpop(self.heap, PointState(dist, p))
            self.furthest_dist = self.heap[0].dist
    
    @property
    def closest(self):
        return [ ps.point for ps in self.heap ]