diff globalneighbors/distance.py @ 5:7e27e874655b

test a larger grid + move distance insertion to its own function
author Jeff Hammel <k0scist@gmail.com>
date Sat, 24 Jun 2017 15:16:10 -0700
parents 49aae0c0293b
children 316e1d54ffd4
line wrap: on
line diff
--- a/globalneighbors/distance.py	Sat Jun 24 14:48:55 2017 -0700
+++ b/globalneighbors/distance.py	Sat Jun 24 15:16:10 2017 -0700
@@ -32,6 +32,23 @@
     return degrees*DEGREESTORADIANS
 
 
+def insert_distance(distances, i, new_distance, k):
+    """
+    insert `(i, new_distance)` into `distances`
+    keeping distances in order with maximum of size `k`
+    """
+
+     # TODO: Binary Search Tree
+    for _index, (geoid, old_distance) in enumerate(distances):
+        if new_distance < old_distance:
+            distances.insert(_index, (i, new_distance))
+            if len(distances) == k+1:
+                distances.pop()
+            break
+    else:
+        distances.append((i, new_distance))
+
+
 def calculate_distances(locations, r=Rearth):
     """
     WARNING! This is an N-squared approach
@@ -101,15 +118,7 @@
                 if len(distances) == k and new_distance >= distances[-1][-1]:
                     continue
 
-                # TODO: Binary Search Tree
-                for _index, (geoid, old_distance) in enumerate(distances):
-                    if new_distance < old_distance:
-                        distances.insert(_index, (i, new_distance))
-                        if len(distances) == k+1:
-                            distances.pop()
-                        break
-                else:
-                    distances.append((i, new_distance))
+                insert_distance(distances, i, new_distance, k)
 
     return neighbors
 
@@ -143,7 +152,7 @@
                                     output=options.output_counter)
 
     # output
-    options.output.write(json.dumps(neighbors, indent=2))
+    options.output.write(json.dumps(neighbors))
 
 if __name__ == '__main__':
     main()