Dijkstra algorithm implementation in Swift

Question

I've implemented Dijkstra's algorithm to find the minimum path between two nodes.

class Node : CustomStringConvertible {
    // unique identifier required for each node
    var identifier : Int
    var distance : Int = Int.max
    var edges = [Edge]()
    var visited = false

    var description: String {
        return "identifier: " + identifier.description + ", Edges: " + edges.count.description
    }

    init(visited: Bool, identifier: Int, edges: [Edge]) {
        self.visited = visited
        self.identifier = identifier
        self.edges = edges
    }

    static func == (lhs: Node, rhs: Node) -> Bool {
        return lhs.identifier == rhs.identifier
    }
}

class Edge {
    var from: Node // does not actually need to be stored!
    var to: Node
    var weight: Int
    var description : String {
        return "from: " + from.description + ", to: " + to.description
    }
    init(to: Node, from: Node, weight: Int) {
        self.to = to
        self.weight = weight
        self.from = from
    }
}

class Graph {
    var nodes: [Node] = []
}


// Complete the quickestWayUp function below.
func setupGraphwith(edges: [[Int]]) -> Graph {
    let graph = Graph()

    // create all the nodes
    // The first and last node need to be included, so need nodes from "to" and "from"
    let nodeNames = Set ( edges.map{ $0[0] } + edges.map{ $0[1]} )
    for node in nodeNames {
        let newNode = Node(visited: false, identifier: node, edges: [])
        graph.nodes.append(newNode)
    }

    // create all the edges to link the nodes
    for edge in edges {
        if let fromNode = graph.nodes.first(where: { $0.identifier == edge[0] }) {
            if let toNode = graph.nodes.first(where: { $0.identifier == edge[1] }) {
                let edge = Edge(to: toNode, from: fromNode, weight: edge[2])
                fromNode.edges.append(edge)
            }
        }
    }
    return graph
}

func shortestPath (source: Int, destination: Int, graph: Graph) -> Int {
    print (graph.nodes)
    var currentNode = graph.nodes.first{ $0.identifier == source }!
    currentNode.visited = true
    currentNode.distance = 0
    var toVisit = [Node]()
    toVisit.append(currentNode)
    while ( !toVisit.isEmpty) {
        toVisit = toVisit.filter{ $0.identifier != currentNode.identifier }
        currentNode.visited = true
        // Go to each adjacent vertex and update the path length
        for connectedEdge in currentNode.edges {
            let dist = currentNode.distance + connectedEdge.weight
            if (dist < connectedEdge.to.distance) {
                connectedEdge.to.distance = dist
                toVisit.append(connectedEdge.to)
                if (connectedEdge.to.visited == true) {
                    connectedEdge.to.visited = false
                }
            }
        }

        currentNode.visited = true

        //set current node to the smallest vertex
        if !toVisit.isEmpty {
            currentNode = toVisit.min(by: { (a, b) -> Bool in
                return a.distance < b.distance
            })!
        }

        if (currentNode.identifier == destination) {
            return currentNode.distance
        }
    }

    return -1
}

class dfsTests: XCTestCase {
    var simpleGraph = Graph()
    override func setUp() {
        simpleGraph = setupGraphwith(edges: [[1,2,4], [1,3,4], [3,4,3], [2,3,2], [3,5,1], [4,6,2], [5,6,3], [3,6,6] ])

    }
    func testbfsReturnVals() {

        XCTAssertEqual(shortestPath(source: 1, destination: 6, graph: simpleGraph), 8 )
    }
}

dfsTests.defaultTestSuite.run()

This is just to practice the implementation of the code, with a view to interviews. Any comments are appreciated.

Answer 1

The algorithm

Let's first have a look at your implementation of Dijkstra's algorithm in the

func shortestPath(source: Int, destination: Int, graph: Graph) -> Int

function:

• The visited property of a node is updated, but nowhere tested. As a consequence, the distance of all neighbors of a node is updated, not only the distance of unvisited neighbors. This does not lead to wrong results (as far as I can see) but causes unnecessary comparisons.

• It is unclear to me what this is for:

          if (connectedEdge.to.visited == true) {
              connectedEdge.to.visited = false
          }
  

  The Dijkstra algorithm does not mark nodes as unvisited, and I cannot see when the condition should be true at all.

• The algorithm does not work if start node and destination node are identical. This can be fixed by computing the next current node (and comparing it with the destination) at the beginning of the main loop, not at the end. This also removes the necessity to check if !toVisit.isEmpty twice (as the loop condition and again inside the loop).

• The distance of the initial node is set to zero twice.

Here is a possible implementation which fixes the above issues:

func shortestPath(source: Int, destination: Int, graph: Graph) -> Int {

    let startNode = graph.nodes.first{ $0.identifier == source }!
    startNode.distance = 0
    var toVisit = [startNode]

    while (!toVisit.isEmpty) {
        // Select node with smallest distance.
        let currentNode = toVisit.min(by: { (a, b) -> Bool in
            return a.distance < b.distance
        })!

        // Destination reached?
        if currentNode.identifier == destination {
            return currentNode.distance
        }

        // Mark as visited.
        currentNode.visited = true
        toVisit = toVisit.filter { $0.identifier != currentNode.identifier }

        // Update unvisited neighbors.
        for edge in currentNode.edges where !edge.to.visited {
            let neighbor = edge.to
            toVisit.append(neighbor)
            let dist = currentNode.distance + edge.weight
            if (dist < neighbor.distance) {
                neighbor.distance = dist
            }
        }

    }

    // Destination not reachable.
    return -1
}

There is one more issue: The toVisit array (which is build “on the fly” while traversing the graph) can contain duplicate elements. A possible fix is to use a set instead of an array. This would required the Node class to be Hashable – see below.

Code and design improvements

• Your function returns -1 if there is no path from the start to the destination node. The Swift way of returning a value or failure is to return an optional, where nil means “no result.”

• All properties of class Edge are never mutated after object creation, they should be declared as constants (with let).

• The Node initializer needs only the identifier – and that should be a constant property.

• It is not possible to compute the shortestPath() function multiple times (with different parameters) because it relies on the visited and distance property to be initialized.

• I would replace the function

  func setupGraphwith(edges: [[Int]]) -> Graph
  

  by an initalizer of the Graph class, and

  func shortestPath(source: Int, destination: Int, graph: Graph) -> Int
  

  by a method of that class:

   class Graph {
  
      init(edgeList: [[Int]]) 
      func distance(from: Int, to: Int) -> Int?
  }
  

  The usage would then be

  let graph = Graph(edgeList: ...)
  let distance = graph.distance(from: 1, to: 6)
  

• The optional bindings in setupGraphwith(edges:)

  if let fromNode = graph.nodes.first(where: { $0.identifier == edge[0] }) 
  if let toNode = graph.nodes.first(where: { $0.identifier == edge[1] })
  

  cannot fail because those nodes were all created before. Therefore a forced unwrap would be appropriate:

  let fromNode = graph.nodes.first(where: { $0.identifier == edge[0] })!
  let toNode = graph.nodes.first(where: { $0.identifier == edge[1] })!
  

  An alternative would be to have a “find or create” method in the Graph class.

Performance improvements

At various points, arrays are used to store, locate, and remove nodes. Each lookup requires a traversal of the array. Using sets would be more efficient. That requires the Node class to be Hashable. I would base equality (and consequently, the hash value) on object identity instead of the numerical identifier.

A priority queue would be more efficient to determine the node with the currently minimum distance.

Putting it all together

Summarizing the above suggestions (with the exception of the priority queue) the code code look like this. I have only omitted the CustomStringConvertible conformance for brevity.

class Node {

    let identifier: Int
    var distance = Int.max
    var edges = [Edge]()
    var visited = false

    init(identifier: Int) {
        self.identifier = identifier
    }
}

extension Node: Hashable {
    static func == (lhs: Node, rhs: Node) -> Bool {
        return lhs === rhs
    }

    func hash(into hasher: inout Hasher) {
        hasher.combine(ObjectIdentifier(self).hashValue)
    }
}

class Edge {
    let from: Node
    let to: Node
    let weight: Int

    init(to: Node, from: Node, weight: Int) {
        self.to = to
        self.from = from
        self.weight = weight
    }
}

class Graph {
    var nodes: Set<Node>

    // Find or create node with the given identifier
    func node(identifier: Int) -> Node {
        if let node = nodes.first(where: { $0.identifier == identifier }) {
            return node
        } else {
            let node = Node(identifier: identifier)
            nodes.insert(node)
            return node
        }
    }

    init(edgeList: [[Int]]) {
        nodes = []
        for edgeDescription in edgeList {
            let fromNode = node(identifier: edgeDescription[0])
            let toNode = node(identifier: edgeDescription[1])
            let edge = Edge(to: toNode, from: fromNode, weight: edgeDescription[2])
            fromNode.edges.append(edge)
        }
    }

    func distance(from: Int, to: Int) -> Int? {
        guard let fromNode = nodes.first(where: { $0.identifier == from }) else {
            return nil
        }
        guard let toNode = nodes.first(where: { $0.identifier == to }) else {
            return nil
        }

        if fromNode == toNode { return 0 }

        for node in nodes {
            node.visited = false
            node.distance = Int.max
        }

        fromNode.distance = 0
        var toVisit = Set([fromNode])

        while !toVisit.isEmpty {
            // Select node with smallest distance.
            let currentNode = toVisit.min(by: { (a, b) -> Bool in
                return a.distance < b.distance
            })!

            // Destination reached?
            if currentNode == toNode { return currentNode.distance }

            // Mark as visited.
            currentNode.visited = true
            toVisit.remove(currentNode)

            // Update unvisited neighbors.
            for edge in currentNode.edges where !edge.to.visited {
                let neighbor = edge.to
                toVisit.insert(neighbor)
                let dist = currentNode.distance + edge.weight
                if (dist < neighbor.distance) {
                    neighbor.distance = dist
                }
            }
        }

        return nil
    }
}