ListofDepths.swift

import PlaygroundSupport

class Node {
  var data: Int
  var next: Node?
  
  init(data: Int) {
    self.data = data
  }
  
  init(treeNode: TreeNode) {
    self.data = treeNode.data
  }
}

class LinkedList: CustomStringConvertible {
  var description: String {
    var result = ""
    var current: Node? = head
    while let node = current {
      result.append("\(node.data) -> ")
      current = current?.next
    }
    result.append("nil")
    return result
  }
  
  var head: Node?
  
  init(head: Node) {
    self.head = head
  }
  
  init(treeNode: TreeNode) {
    self.head = Node(treeNode: treeNode)
  }
  
  func append(node: Node) {
    guard let head = head else {
      self.head = node
      return
    }
    
    var current: Node? = head
    while (current?.next != nil) {
      current = current?.next
    }
    current?.next = node
  }
}

class TreeNode {
  var data: Int
  var left: TreeNode?
  var right: TreeNode?
  
  init(data: Int) {
    self.data = data
  }
}

class Tree {
  var root: TreeNode
  
  init(root: TreeNode) {
    self.root = root
  }
  
  func printLinkedListsUsingBFS() {
    for list in linkedListsUsingBFS() {
      print(list)
    }
  }
  
  func printLinkedListsUsingDFS() {
    for list in linkedListsUsingDFS() {
      print(list)
    }
  }
  
  func linkedListsUsingBFS() -> [LinkedList] {
    var lists = [LinkedList]()
    var parents = [root]
    var children = [TreeNode]()
    
    var level = 0
    while !parents.isEmpty {
      children = []
      for parent in parents {
        if let left = parent.left {
          children.append(left)
        }
        if let right = parent.right {
          children.append(right)
        }
        link(treeNode: parent, to: &lists, at: level)
      }
      parents = children
      level += 1
    }
    return lists
  }
  
  func linkedListsUsingDFS() -> [LinkedList] {
    var lists = [LinkedList]()
    linkUsingDFS(nodes: [root], to: &lists, at: 0)
    return lists
  }
  
  func linkUsingDFS(nodes: [TreeNode], to lists: inout [LinkedList], at level: Int) {
    for node in nodes {
      link(treeNode: node, to: &lists, at: level)
      var children = [TreeNode]()
      if let left = node.left {
        children.append(left)
      }
      if let right = node.right {
        children.append(right)
      }
      linkUsingDFS(nodes: children, to: &lists, at: level+1)
    }
  }
  
  func link(treeNode: TreeNode, to lists: inout [LinkedList], at level: Int) {
    let list: LinkedList
    if lists.count > level {
      list = lists[level]
      list.append(node: Node(treeNode: treeNode))
    } else {
      list = LinkedList(treeNode: treeNode)
      lists.append(list)
    }
  }
}

let node1 = TreeNode(data: 1)
let node2 = TreeNode(data: 2)
let node3 = TreeNode(data: 3)
let node4 = TreeNode(data: 4)
let node5 = TreeNode(data: 5)
let node6 = TreeNode(data: 6)
let node7 = TreeNode(data: 7)
let node8 = TreeNode(data: 8)
let node9 = TreeNode(data: 9)
let node10 = TreeNode(data: 10)
let node11 = TreeNode(data: 11)

node1.left = node2
node1.right = node3

node2.left = node4
node2.right = node5

node3.left = node6
node3.right = node7

node4.left = node8
node4.right = node9

node5.left = node10
node5.right = node11

let tree = Tree(root: node1)

print("List of depths using BFS: ")
tree.printLinkedListsUsingBFS()

print("\nList of depths using DFS: ")
tree.printLinkedListsUsingDFS()