Mercurial > hg > GlobalNeighbors
diff globalneighbors/distance.py @ 8:e3d6919130ca
insert via BST
author | Jeff Hammel <k0scist@gmail.com> |
---|---|
date | Sun, 25 Jun 2017 11:21:28 -0700 |
parents | 254195d0bac2 |
children | 638fad06e556 |
line wrap: on
line diff
--- a/globalneighbors/distance.py Sun Jun 25 09:13:48 2017 -0700 +++ b/globalneighbors/distance.py Sun Jun 25 11:21:28 2017 -0700 @@ -3,6 +3,7 @@ """ import argparse +import bisect import json import sys import time @@ -51,8 +52,19 @@ else: distances.append((i, new_distance)) +class KeyWrapper: + def __init__(self, iterable, key): + self.it = iterable + self.key = key + + def __getitem__(self, i): + return self.key(self.it[i]) + + def __len__(self): + return len(self.it) def insert_distance_bisect(distances, i, new_distance, k): + if not distances: distances.append((i, new_distance)) return @@ -60,10 +72,13 @@ if len(distances) < k: distances.append((i, new_distance)) return - indices = [0, len(distances)] - while True: - midpoint = int((indices[-1] - indices[0])/2) + point = bisect.bisect_left(KeyWrapper(distances, + key=lambda x: x[-1]), + new_distance) + distances.insert(point, (i, new_distance)) + if len(distances) == k+1: + distances.pop() def calculate_distances(locations, r=Rearth): """