Skip to content

Instantly share code, notes, and snippets.

@ageron
Last active May 5, 2020 01:43
Show Gist options
  • Save ageron/b31a5489fc84c96184187a171a4923e1 to your computer and use it in GitHub Desktop.
Save ageron/b31a5489fc84c96184187a171a4923e1 to your computer and use it in GitHub Desktop.
public class Vertex {
var key: String
public init(key: String) {
self.key = key
print("Creating vertex \(key)")
}
deinit {
print("Removing vertex \(key)")
}
}
public class Edge {
let source: Vertex
let target: Vertex
var weight: Int
init(from source: Vertex, to target: Vertex, weight: Int = 0) {
self.source = source
self.target = target
self.weight = weight
print("Creating edge \(source.key) -> \(target.key) (weight \(weight))")
}
deinit {
print("Deleting edge \(source.key) -> \(target.key) (weight \(weight))")
}
}
public protocol Graph {
func addEdge(_ edge: Edge)
func edges(of vertex: Vertex) -> [Edge]
func neighbors(of vertex: Vertex) -> [Vertex]
}
public class UndirectedGraph: Graph {
private var index = Dictionary<String, Array<Edge>>()
public func addEdge(_ edge: Edge) {
index[edge.source.key, default: []].append(edge)
if edge.source.key != edge.target.key {
index[edge.target.key, default: []].append(edge)
}
}
public func edges(of vertex: Vertex) -> [Edge] {
return index[vertex.key, default: []]
}
public func neighbors(of vertex: Vertex) -> [Vertex] {
return edges(of: vertex).map {
$0.source.key == vertex.key ? $0.target : $0.source
}
}
}
public class DirectedGraph: Graph {
private var sourceIndex = Dictionary<String, Array<Edge>>()
private var targetIndex = Dictionary<String, Array<Edge>>()
public func addEdge(_ edge: Edge) {
sourceIndex[edge.source.key, default: []].append(edge)
targetIndex[edge.target.key, default: []].append(edge)
}
public func edges(source: Vertex) -> [Edge] {
return sourceIndex[source.key, default: []]
}
public func edges(target: Vertex) -> [Edge] {
return targetIndex[target.key, default: []]
}
public func edges(of vertex: Vertex) -> [Edge] {
return edges(source: vertex) + edges(target: vertex)
}
public func targets(of vertex: Vertex) -> [Vertex] {
return edges(source: vertex).map { $0.target }
}
public func sources(of vertex: Vertex) -> [Vertex] {
return edges(target: vertex).map { $0.source }
}
public func neighbors(of vertex: Vertex) -> [Vertex] {
return sources(of: vertex) + targets(of: vertex)
}
}
func testGraph(_ graph: Graph) {
let a = Vertex(key: "a")
let b = Vertex(key: "b")
let c = Vertex(key: "c")
let d = Vertex(key: "d")
graph.addEdge(Edge(from: a, to: b))
graph.addEdge(Edge(from: b, to: c))
graph.addEdge(Edge(from: c, to: a))
graph.addEdge(Edge(from: a, to: d))
graph.addEdge(Edge(from: a, to: a))
if let directedGraph = graph as? DirectedGraph {
print("Targets of vertex a:")
directedGraph.targets(of: a).forEach { print($0.key) }
print("Sources of vertex a:")
directedGraph.sources(of: a).forEach { print($0.key) }
} else {
print("Neighbors of vertex a:")
graph.neighbors(of: a).map { print($0.key) }
}
}
testGraph(DirectedGraph())
testGraph(UndirectedGraph())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment