|
| 1 | +// Written by Alejandro Isaza. |
| 2 | + |
| 3 | +import Foundation |
| 4 | + |
| 5 | +public protocol Graph { |
| 6 | + associatedtype Vertex: Hashable |
| 7 | + associatedtype Edge: WeightedEdge where Edge.Vertex == Vertex |
| 8 | + |
| 9 | + /// Lists all edges going out from a vertex. |
| 10 | + func edgesOutgoing(from vertex: Vertex) -> [Edge] |
| 11 | +} |
| 12 | + |
| 13 | +public protocol WeightedEdge { |
| 14 | + associatedtype Vertex |
| 15 | + |
| 16 | + /// The edge's cost. |
| 17 | + var cost: Double { get } |
| 18 | + |
| 19 | + /// The target vertex. |
| 20 | + var target: Vertex { get } |
| 21 | +} |
| 22 | + |
| 23 | +public final class AStar<G: Graph> { |
| 24 | + /// The graph to search on. |
| 25 | + public let graph: G |
| 26 | + |
| 27 | + /// The heuristic cost function that estimates the cost between two vertices. |
| 28 | + /// |
| 29 | + /// - Note: The heuristic function needs to always return a value that is lower-than or equal to the actual |
| 30 | + /// cost for the resulting path of the A* search to be optimal. |
| 31 | + public let heuristic: (G.Vertex, G.Vertex) -> Double |
| 32 | + |
| 33 | + /// Open list of nodes to expand. |
| 34 | + private var open: HashedHeap<Node<G.Vertex>> |
| 35 | + |
| 36 | + /// Closed list of vertices already expanded. |
| 37 | + private var closed = Set<G.Vertex>() |
| 38 | + |
| 39 | + /// Actual vertex cost for vertices we already encountered (refered to as `g` on the literature). |
| 40 | + private var costs = Dictionary<G.Vertex, Double>() |
| 41 | + |
| 42 | + /// Store the previous node for each expanded node to recreate the path. |
| 43 | + private var parents = Dictionary<G.Vertex, G.Vertex>() |
| 44 | + |
| 45 | + /// Initializes `AStar` with a graph and a heuristic cost function. |
| 46 | + public init(graph: G, heuristic: @escaping (G.Vertex, G.Vertex) -> Double) { |
| 47 | + self.graph = graph |
| 48 | + self.heuristic = heuristic |
| 49 | + open = HashedHeap(sort: <) |
| 50 | + } |
| 51 | + |
| 52 | + /// Finds an optimal path between `source` and `target`. |
| 53 | + /// |
| 54 | + /// - Precondition: both `source` and `target` belong to `graph`. |
| 55 | + public func path(start: G.Vertex, target: G.Vertex) -> [G.Vertex] { |
| 56 | + open.insert(Node<G.Vertex>(vertex: start, cost: 0, estimate: heuristic(start, target))) |
| 57 | + while !open.isEmpty { |
| 58 | + guard let node = open.remove() else { |
| 59 | + break |
| 60 | + } |
| 61 | + costs[node.vertex] = node.cost |
| 62 | + |
| 63 | + if (node.vertex == target) { |
| 64 | + let path = buildPath(start: start, target: target) |
| 65 | + cleanup() |
| 66 | + return path |
| 67 | + } |
| 68 | + |
| 69 | + if !closed.contains(node.vertex) { |
| 70 | + expand(node: node, target: target) |
| 71 | + closed.insert(node.vertex) |
| 72 | + } |
| 73 | + } |
| 74 | + |
| 75 | + // No path found |
| 76 | + return [] |
| 77 | + } |
| 78 | + |
| 79 | + private func expand(node: Node<G.Vertex>, target: G.Vertex) { |
| 80 | + let edges = graph.edgesOutgoing(from: node.vertex) |
| 81 | + for edge in edges { |
| 82 | + let g = cost(node.vertex) + edge.cost |
| 83 | + if g < cost(edge.target) { |
| 84 | + open.insert(Node<G.Vertex>(vertex: edge.target, cost: g, estimate: heuristic(edge.target, target))) |
| 85 | + parents[edge.target] = node.vertex |
| 86 | + } |
| 87 | + } |
| 88 | + } |
| 89 | + |
| 90 | + private func cost(_ vertex: G.Edge.Vertex) -> Double { |
| 91 | + if let c = costs[vertex] { |
| 92 | + return c |
| 93 | + } |
| 94 | + |
| 95 | + let node = Node(vertex: vertex, cost: Double.greatestFiniteMagnitude, estimate: 0) |
| 96 | + if let index = open.index(of: node) { |
| 97 | + return open[index].cost |
| 98 | + } |
| 99 | + |
| 100 | + return Double.greatestFiniteMagnitude |
| 101 | + } |
| 102 | + |
| 103 | + private func buildPath(start: G.Vertex, target: G.Vertex) -> [G.Vertex] { |
| 104 | + var path = Array<G.Vertex>() |
| 105 | + path.append(target) |
| 106 | + |
| 107 | + var current = target |
| 108 | + while current != start { |
| 109 | + guard let parent = parents[current] else { |
| 110 | + return [] // no path found |
| 111 | + } |
| 112 | + current = parent |
| 113 | + path.append(current) |
| 114 | + } |
| 115 | + |
| 116 | + return path.reversed() |
| 117 | + } |
| 118 | + |
| 119 | + private func cleanup() { |
| 120 | + open.removeAll() |
| 121 | + closed.removeAll() |
| 122 | + parents.removeAll() |
| 123 | + } |
| 124 | +} |
| 125 | + |
| 126 | +private struct Node<V: Hashable>: Hashable, Comparable { |
| 127 | + /// The graph vertex. |
| 128 | + var vertex: V |
| 129 | + |
| 130 | + /// The actual cost between the start vertex and this vertex. |
| 131 | + var cost: Double |
| 132 | + |
| 133 | + /// Estimated (heuristic) cost betweent this vertex and the target vertex. |
| 134 | + var estimate: Double |
| 135 | + |
| 136 | + public init(vertex: V, cost: Double, estimate: Double) { |
| 137 | + self.vertex = vertex |
| 138 | + self.cost = cost |
| 139 | + self.estimate = estimate |
| 140 | + } |
| 141 | + |
| 142 | + static func < (lhs: Node<V>, rhs: Node<V>) -> Bool { |
| 143 | + return lhs.cost + lhs.estimate < rhs.cost + rhs.estimate |
| 144 | + } |
| 145 | + |
| 146 | + static func == (lhs: Node<V>, rhs: Node<V>) -> Bool { |
| 147 | + return lhs.vertex == rhs.vertex |
| 148 | + } |
| 149 | + |
| 150 | + var hashValue: Int { |
| 151 | + return vertex.hashValue |
| 152 | + } |
| 153 | +} |
0 commit comments