Recurse Center Interview

DepthFirstSearch.swift

import UIKit
import Foundation

extension Int {
    func times(_ perform: (Int) -> Void) {
        (0..<self).forEach(perform)
    }
}

class Node<T> {
    var parent: Node?
    var children: [Node<T>] = []
    let value: T
    let isRoot: Bool
    var isVisited: Bool = false
    
    init(value: T, isRoot: Bool? = nil) {
        self.value = value
        self.isRoot = isRoot ?? false
    }
    
    func addChild(_ node: Node<T>) {
        node.parent = self
        children.append(node)
    }
    
    var level: Int {
        
        guard let parent = parent else {
            return 0
        }
        
        return parent.level + 1
    }
}

extension Node where T: Equatable {
    
    // Don't necessarily need this "visited" marker for DFS in a tree
    func depthFirstSearch(_ value: T) -> Node? {
        
        // Mark current node as visited
        isVisited = true
        
        // Check it it's the current node's value
        guard value != self.value else {
            return self
        }
        
        // Go down one child, and mark it as visited
        for child in children {
            if !child.isVisited {
                child.isVisited = true
                if let found = child.depthFirstSearch(value) {
                    return found
                }
            }
        }
        
        return nil
    }
}

extension Node: CustomStringConvertible {
    var description: String {
        
        var s = "\(value)"
        
        if !children.isEmpty {
            var tabs = "\t"
            level.times{ _ in tabs.append("\t") }
            s += " \n\(tabs)" + children.map{ $0.description }.joined(separator: "\n " + tabs)
        }
        return s
    }
}

let tree = Node<String>(value: "a", isRoot: true)

// Add to root
let bNode = Node<String>(value: "b")
let cNode = Node<String>(value: "c")
tree.addChild(bNode)
tree.addChild(cNode)

// Add to b
let dNode = Node<String>(value: "d")
let eNode = Node<String>(value: "e")
let fNode = Node<String>(value: "f")
bNode.addChild(dNode)
bNode.addChild(eNode)
bNode.addChild(fNode)

// Add to c
let gNode = Node<String>(value: "g")
cNode.addChild(gNode)

// Add to g
let hNode = Node<String>(value: "h")
gNode.addChild(hNode)

//print(tree.description)

if let g = tree.depthFirstSearch("g") {
    print(g.value)
} else {
    print("Couldn't find g node :(")
}