Graph + DepthFirstSearch

DepthFirstSearch.swift

import Foundation

final class DepthFirstSearch<T: Hashable> {

    private let graph: Graph<T>
    private var searchStack: [T] = []
    private var visited: Set<T> = Set()

    init(graph: Graph<T>) {
        self.graph = graph
    }

    func contains(_ value: T) -> Bool {
        searchStack.append(graph.storage.first!.key)

        while let vertex = searchStack.popLast() {
            if vertex == value { return true }

            if !visited.contains(vertex) {
                searchStack.append(contentsOf: graph.getEdges(for: vertex) ?? [])
                visited.insert(vertex)
            } else {
                continue
            }
        }

        return false
    }
}

final class Graph<T: Hashable> {

    private(set) var storage: [T: Set<T>]

    init(vertices: [T]) {
        storage = [:]
        vertices.forEach { storage[$0] = Set() }
    }

    func addEdge(between v: T, and w: T) {
        guard storage[v] != nil && storage[w] != nil else { return }
        storage[v]?.insert(w)
        storage[w]?.insert(v)
    }

    func addEdges(_ edges: [T], to vertex: T) {
        edges.forEach { addEdge(between: vertex, and: $0) }
    }

    func getEdges(for vertex: T) -> Set<T>? {
        return storage[vertex]
    }
}