Trees and Graphs

Setup

This chapter uses the TreeNode class from the Building blocks chapter:

final class TreeNode {
    var val: Int
    var left: TreeNode?
    var right: TreeNode?

    init(_ val: Int, _ left: TreeNode? = nil, _ right: TreeNode? = nil) {
        self.val = val
        self.left = left
        self.right = right
    }
}

Tree DFS, recursive

DFS (depth-first search) goes as deep as possible before it comes back. On trees, recursion is the shortest way to write it.

Return a value up from the children:

func maxDepth(_ root: TreeNode?) -> Int {
    guard let root else { return 0 }
    return 1 + max(maxDepth(root.left), maxDepth(root.right))
}

Collect values with a nested function. It can change variables of the outer function:

func inorder(_ root: TreeNode?) -> [Int] {
    var result: [Int] = []
    func visit(_ node: TreeNode?) {
        guard let node else { return }
        visit(node.left)
        result.append(node.val)   // pre-order: do this first; post-order: do this last
        visit(node.right)
    }
    visit(root)
    return result
}

Pass limits down to the children. Here each node must stay between a low and a high value:

func isValidBST(_ root: TreeNode?) -> Bool {
    func check(_ node: TreeNode?, low: Int?, high: Int?) -> Bool {
        guard let node else { return true }
        if let low, node.val <= low { return false }
        if let high, node.val >= high { return false }
        return check(node.left, low: low, high: node.val)
            && check(node.right, low: node.val, high: high)
    }
    return check(root, low: nil, high: nil)
}

Return two things at once with an optional or a tuple. Here nil means "not balanced":

func isBalanced(_ root: TreeNode?) -> Bool {
    func height(_ node: TreeNode?) -> Int? {
        guard let node else { return 0 }
        guard let left = height(node.left),
              let right = height(node.right),
              abs(left - right) <= 1 else { return nil }
        return 1 + max(left, right)
    }
    return height(root) != nil
}

let root = TreeNode(4, TreeNode(2, TreeNode(1), TreeNode(3)), TreeNode(7))
print(maxDepth(root), inorder(root), isValidBST(root), isBalanced(root))
// 3 [1, 2, 3, 4, 7] true true

Tree DFS, iterative

Recursion uses the call stack. A very deep tree (for example 100,000 nodes in a line) can crash with a stack overflow. An explicit stack avoids that.

func preorderIterative(_ root: TreeNode?) -> [Int] {
    guard let root else { return [] }
    var result: [Int] = []
    var stack = [root]
    while let node = stack.popLast() {
        result.append(node.val)
        if let right = node.right { stack.append(right) }   // push right first,
        if let left = node.left { stack.append(left) }      // so left comes out first
    }
    return result
}

func inorderIterative(_ root: TreeNode?) -> [Int] {
    var result: [Int] = []
    var stack: [TreeNode] = []
    var current = root
    while current != nil || !stack.isEmpty {
        while let node = current {        // go left as far as possible
            stack.append(node)
            current = node.left
        }
        let node = stack.removeLast()
        result.append(node.val)
        current = node.right
    }
    return result
}

Tree BFS (level order)

BFS (breadth-first search) visits nodes level by level. In Swift you can skip the queue: keep the current level in an array and build the next one.

func levelOrder(_ root: TreeNode?) -> [[Int]] {
    guard let root else { return [] }
    var result: [[Int]] = []
    var level = [root]
    while !level.isEmpty {
        result.append(level.map(\.val))
        level = level.flatMap { [$0.left, $0.right].compactMap { $0 } }
    }
    return result
}

let tree = TreeNode(4, TreeNode(2, TreeNode(1), TreeNode(3)), TreeNode(7))
print(levelOrder(tree))   // [[4], [2, 7], [1, 3]]

Use BFS for "level by level" questions, the right side view, and the minimum depth.

Graph DFS

Graphs can have cycles, so you need a visited list. Here: count the connected parts of an undirected graph.

func countComponents(_ n: Int, _ edges: [[Int]]) -> Int {
    var graph = Array(repeating: [Int](), count: n)
    for edge in edges {
        graph[edge[0]].append(edge[1])
        graph[edge[1]].append(edge[0])
    }
    var visited = Array(repeating: false, count: n)

    func dfs(_ node: Int) {
        visited[node] = true
        for next in graph[node] where !visited[next] {
            dfs(next)
        }
    }

    var components = 0
    for node in 0..<n where !visited[node] {
        dfs(node)
        components += 1
    }
    return components
}

print(countComponents(5, [[0, 1], [1, 2], [3, 4]]))   // 2

The iterative version uses a stack. Mark a node as visited when you push it, so it is never pushed twice:

func reachable(from start: Int, in graph: [[Int]]) -> Set<Int> {
    var visited: Set<Int> = [start]
    var stack = [start]
    while let node = stack.popLast() {
        for next in graph[node] where !visited.contains(next) {
            visited.insert(next)
            stack.append(next)
        }
    }
    return visited
}

print(reachable(from: 0, in: [[1], [2], [], [0]]).sorted())   // [0, 1, 2]

Graph BFS: shortest path

In a graph without weights, BFS finds the shortest path. It reaches every node first by the fewest edges.

func shortestPath(from start: Int, to goal: Int, in graph: [[Int]]) -> Int? {
    var distance = Array(repeating: -1, count: graph.count)   // -1 means "not seen"
    distance[start] = 0
    var queue = [start]
    var head = 0
    while head < queue.count {
        let node = queue[head]
        head += 1
        if node == goal { return distance[node] }
        for next in graph[node] where distance[next] == -1 {
            distance[next] = distance[node] + 1
            queue.append(next)
        }
    }
    return nil
}

let graph = [[1, 2], [3], [3], [4], []]
print(shortestPath(from: 0, to: 4, in: graph)!)   // 3: 0 -> 1 -> 3 -> 4

The head index makes the queue O(1) per step (see Building blocks). removeFirst() would make it O(n).

For graphs with weights, use Dijkstra's algorithm: the same idea, but with the heap from Building blocks instead of a queue.

Grids

A grid is a graph. Each cell connects to its 4 neighbors. "Number of islands" is the classic:

func numIslands(_ grid: [[Character]]) -> Int {
    guard let cols = grid.first?.count else { return 0 }
    let rows = grid.count
    var seen = Array(repeating: Array(repeating: false, count: cols), count: rows)
    let directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
    var islands = 0

    for r in 0..<rows {
        for c in 0..<cols where grid[r][c] == "1" && !seen[r][c] {
            islands += 1
            seen[r][c] = true
            var queue = [(r, c)]
            var head = 0
            while head < queue.count {
                let (row, col) = queue[head]
                head += 1
                for (dr, dc) in directions {
                    let nr = row + dr
                    let nc = col + dc
                    guard nr >= 0, nr < rows, nc >= 0, nc < cols,
                          grid[nr][nc] == "1", !seen[nr][nc] else { continue }
                    seen[nr][nc] = true
                    queue.append((nr, nc))
                }
            }
        }
    }
    return islands
}

let grid: [[Character]] = [
    ["1", "1", "0", "0"],
    ["1", "0", "0", "1"],
    ["0", "0", "1", "1"],
]
print(numIslands(grid))   // 2

Tip: the tuple (r, c) cannot go in a Set, so a 2D [[Bool]] array is the easiest "visited" list for grids.

Topological sort

A topological order lists the nodes of a directed graph so that every edge goes from earlier to later. Example: take each course after its prerequisites.

Kahn's algorithm (BFS): start with nodes that have no incoming edges. Remove them, and repeat.

// prerequisites [a, b] means "take b before a".
// Returns an order, or nil if there is a cycle.
func courseOrder(_ n: Int, _ prerequisites: [[Int]]) -> [Int]? {
    var graph = Array(repeating: [Int](), count: n)
    var indegree = Array(repeating: 0, count: n)     // number of incoming edges
    for pair in prerequisites {
        graph[pair[1]].append(pair[0])
        indegree[pair[0]] += 1
    }

    var order = (0..<n).filter { indegree[$0] == 0 }  // no prerequisites
    var head = 0
    while head < order.count {
        let course = order[head]
        head += 1
        for next in graph[course] {
            indegree[next] -= 1
            if indegree[next] == 0 { order.append(next) }
        }
    }
    return order.count == n ? order : nil            // some courses never freed: a cycle
}

print(courseOrder(4, [[1, 0], [2, 0], [3, 1], [3, 2]])!)   // [0, 1, 2, 3]
print(courseOrder(2, [[0, 1], [1, 0]]) as Any)             // nil

The queue array is also the answer. Nice and short.

Cycle check with DFS uses three states: not visited, on the current path, and done. Meeting a node that is "on the current path" means a cycle.

func hasCycle(_ graph: [[Int]]) -> Bool {
    var state = Array(repeating: 0, count: graph.count)   // 0 new, 1 on path, 2 done

    func visit(_ node: Int) -> Bool {
        if state[node] == 1 { return true }    // back on our own path: a cycle
        if state[node] == 2 { return false }
        state[node] = 1
        for next in graph[node] {
            if visit(next) { return true }
        }
        state[node] = 2
        return false
    }

    for node in graph.indices where state[node] == 0 {
        if visit(node) { return true }
    }
    return false
}

print(hasCycle([[1], [2], [0]]))   // true
print(hasCycle([[1], [2], []]))    // false

Union-find

Union-find (also called disjoint set union) keeps groups of nodes. It answers "are these two in the same group?" almost in O(1).

struct UnionFind {
    private var parent: [Int]
    private var rank: [Int]
    private(set) var groups: Int

    init(_ n: Int) {
        parent = Array(0..<n)
        rank = Array(repeating: 0, count: n)
        groups = n
    }

    mutating func find(_ x: Int) -> Int {
        var root = x
        while parent[root] != root { root = parent[root] }
        var node = x
        while parent[node] != root {          // path compression:
            let next = parent[node]           // point every node on the way
            parent[node] = root               // straight to the root
            node = next
        }
        return root
    }

    @discardableResult
    mutating func union(_ a: Int, _ b: Int) -> Bool {
        let rootA = find(a)
        let rootB = find(b)
        if rootA == rootB { return false }    // already in the same group
        if rank[rootA] < rank[rootB] {
            parent[rootA] = rootB
        } else if rank[rootA] > rank[rootB] {
            parent[rootB] = rootA
        } else {
            parent[rootB] = rootA
            rank[rootA] += 1
        }
        groups -= 1
        return true
    }
}

Two tricks make it fast:

Path compression: after a find, nodes point straight to the root.
Union by rank: attach the shorter tree under the taller one.

Use it to find the edge that creates a cycle:

// Nodes are 1...n. Return the first edge that connects two nodes already in one group.
func redundantConnection(_ edges: [[Int]]) -> [Int]? {
    var uf = UnionFind(edges.count + 1)
    for edge in edges where !uf.union(edge[0], edge[1]) {
        return edge
    }
    return nil
}

print(redundantConnection([[1, 2], [1, 3], [2, 3]])!)   // [2, 3]

Which one to use

Question	Tool
Depth, paths, "check every node" on a tree	Recursive DFS
Level by level, minimum depth	BFS
Shortest path, no weights	BFS
Shortest path, with weights	Dijkstra with a heap
Connected parts, islands	DFS, BFS or union-find
Order with dependencies, cycle in a directed graph	Topological sort
Edges arrive one by one, "are a and b connected?"	Union-find

Complexity for DFS, BFS and topological sort: O(V + E) time and O(V) space, where V is the number of nodes and E the number of edges.

#Setup

#Tree DFS, recursive

#Tree DFS, iterative

#Tree BFS (level order)

#Graph DFS

#Graph BFS: shortest path

#Grids

#Topological sort

#Union-find

#Which one to use