pathfinder.py

import csv
import heapq
import math

class Skyland(object):
  def __init__(self, name, location):
    self.name = name
    self.location = location
    self.neighbours = []
  
  def __repr__(self):
    return "Skyland(%r, %r)" % (self.name, self.location)


MAX_EDGE_DISTANCE = 6000

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

def load_skylands(fh):
  r = csv.reader(fh)
  skylands = {}
  for name, x, y in r:
    skylands[name] = Skyland(name, (int(x), int(y)))
  for src in skylands.values():
    for dest in skylands.values():
      if src == dest: continue
      dist = distance(src.location, dest.location)
      if dist <= MAX_EDGE_DISTANCE:
        src.neighbours.append((dist, dest))
    src.neighbours.sort()
  return skylands

def dijkstras(src, dest, max_dist):
  final = {}  # Maps vertices to (min cost, real cost, previous vertex) tuples
  intermediate = {}  # Maps vertices to (min cost, real cost, previous vertex) tuples
  frontier = []  # Heap of (vertex, min cost, real cost) tuples
  
  crow_distance = distance(src.location, dest.location)
  heapq.heappush(frontier, (src, crow_distance, 0))
  intermediate[src] = (crow_distance, 0, None)
  
  while frontier:
    current, min_cost, real_cost = heapq.heappop(frontier)
    if current in final: continue
    
    final[current] = intermediate[current]
    del intermediate[current]
    if current == dest: break
    
    for edge_dist, edge_dest in current.neighbours:
      if edge_dest in final: continue
      if edge_dist > max_dist: break
      
      new_real_cost = real_cost + edge_dist
      new_min_cost = new_real_cost + distance(edge_dest.location, dest.location)
      if edge_dest not in intermediate or intermediate[edge_dest][0] > new_min_cost:
        intermediate[edge_dest] = (new_min_cost, new_real_cost, current)
        heapq.heappush(frontier, (edge_dest, new_min_cost, new_real_cost))
    
  if dest not in final:
    return None, []
  path = []
  current = dest
  while current:
    path.append(current)
    current = final[current][2]
  path.reverse()
  return final[dest][1], path

def all_points(graph, max_dist):
  paths = dict(((src, dest), (dist, None)) for src in graph
               for dist, dest in src.neighbours if dist <= max_dist)
  for via in graph:
    for src in graph:
      for dest in graph:
        a = paths.get((src, via))
        b = paths.get((via, dest))
        c = paths.get((src, dest))
        if a and b and (not c or a[0] + b[0] < c[0]):
          paths[(src, dest)] = (a[0] + b[0], via)
  return paths