A generic AVL tree with SequenceType, CollectionType and ArrayLiteralConvertible extensions

Question

I'm learning Swift and trying get a good understanding of some of its many features. The following is an AVL tree implementation with extensions that conform to SequenceType, CollectionType and ArrayLiteralConvertible

Everything seems to work correctly. I wish to get some feedback on my overall approach, pointers on how I can improve it and conventions I should be following.

protocol BinarySearchTreeType {
  associatedtype Element
  mutating func insert(element: Element)
  mutating func find(element: Element) -> Element?
}


class Node<Element: Comparable> {
  var value: Element
  var leftNode: Node?
  var rightNode: Node?

  init(value: Element) {
    self.value = value
  }

  var height: Int {
    let left = leftNode != nil ? leftNode!.height + 1 : 0
    let right = rightNode != nil ? rightNode!.height + 1: 0
    return max(left, right)
  }
}

enum BinarySearchTreeError: ErrorType {
  case OutOfBound
}

final class BinarySearchTree<Element: Comparable>: BinarySearchTreeType {

  private var _root: Node<Element>?
  private var _count: Int = 0

  var root: Node<Element>? {
    return _root
  }
  var count: Int {
    return _count
  }

  /**
   Inserts an element into the AVL tree. Duplicate elements are ignored.
   - parameter element: Element which will be added to the tree.
   It must conform to Comparable.
   */
  func insert (element: Element) {
    if let _ = _root {
      self.insert(element, currentNode: &_root!)
    } else {
      _root = Node(value: element)
      _count += 1
    }
  }

  private func insert(element: Element, inout currentNode: Node<Element>) -> Node<Element> {
    if currentNode.value > element {
      if currentNode.leftNode != nil {
        currentNode.leftNode = insert(element, currentNode: &currentNode.leftNode!)
      } else {
        currentNode.leftNode = Node<Element>(value: element)
        _count += 1
      }
      if height(currentNode.leftNode) - height(currentNode.rightNode) == 2 {
        if element < currentNode.leftNode!.value {
          currentNode = rightRotate(currentNode)
        } else {
          currentNode.leftNode = leftRotate(currentNode.leftNode!)
          currentNode = rightRotate(currentNode)
        }
      }
    } else if currentNode.value < element {
      if currentNode.rightNode != nil {
        currentNode.rightNode = insert(element, currentNode: &currentNode.rightNode!)
      } else {
        currentNode.rightNode = Node<Element>(value: element)
        _count += 1
      }

      if height(currentNode.rightNode) - height(currentNode.leftNode) == 2 {
        if element > currentNode.rightNode!.value {
          currentNode = leftRotate(currentNode)
        } else {
          currentNode.rightNode = rightRotate(currentNode.rightNode!)
          currentNode = leftRotate(currentNode)
        }
      }
    }
    return currentNode
  }

  private func balance(node: Node<Element>?) -> Int {
    if let node = node {
      return height(node.leftNode) - height(node.rightNode)
    }
    return 0
  }

  private func height (node: Node<Element>?) -> Int {
    return node != nil ? node!.height : -1
  }

  private func leftRotate(node: Node<Element>) -> Node<Element> {
    let newRoot = node.rightNode!
    let oldRootRight = newRoot.leftNode
    node.rightNode = oldRootRight
    newRoot.leftNode = node
    return newRoot
  }

  private func rightRotate(node: Node<Element>) -> Node<Element> {
    let newRoot = node.leftNode!
    let oldRootLeft = newRoot.rightNode
    node.leftNode = oldRootLeft
    newRoot.rightNode = node
    return newRoot
  }

  /**
   Returns element with the smallest value in the tree.
   - parameter root: Root node for the tree.
   - returns: Node with the smallest value
   */
  func minNode (root: Node<Element>) -> Node<Element> {
    var _current = root
    while _current.leftNode != nil {
      _current = _current.leftNode!
    }
    return _current
  }

  /**
   Returns element with the largest value in the tree.
   - parameter root: Root node for the tree.
   - returns: Node with the largest value
   */
  func maxNode(root: Node<Element>) -> Node<Element> {
      var _current = root
      while _current.rightNode != nil {
        _current = _current.rightNode!
      }
      return _current
  }
  /**
   Returns the predecessor element of a given element according to in-order traversal of the tree.
   - parameter node: Node from which want the predecessor.
   - returns: Predecessor for the node passed as a parameter
   */
  func predecessor(node: Node<Element>) -> Node<Element>? {
    if let root = _root {
      return self.predecessor(node, root: root)
    }
    return nil
  }

  private func predecessor(node: Node<Element>, root: Node<Element>) -> Node<Element>? {
    if let leftSubTree = node.leftNode {
      return maxNode(leftSubTree)
    }

    var _current: Node<Element>? = root
    var result: Node<Element>?
    while _current != nil {
      if _current!.value > node.value {
        _current = _current?.leftNode
      } else if _current!.value < node.value {
        result = _current
        _current = _current!.rightNode
      } else {
        return result
      }
    }
    return result
  }

  /**
   Returns the successor element of a given element according to in-order traversal of the tree.
   - parameter node: Node from which we want the successor.
   - returns: Predecessor for the node passed as a parameter
   */
  func successor(node: Node<Element>) -> Node<Element>? {
    if let root = _root {
      return successor(node, root: root)
    }
    return nil
  }

  private func successor(node: Node<Element>, root: Node<Element>) -> Node<Element>? {
    if let righSubTree = node.rightNode {
      return minNode(righSubTree)
    }

    var _current: Node<Element>? = root
    var result: Node<Element>?

    while _current != nil {
      if node.value < _current!.value {
        result = _current
        _current = _current?.leftNode
      } else if node.value > _current!.value {
        _current = _current?.rightNode
      } else {
        return result
      }
    }
    return result
  }
  /**
   Finds an element in the AVL tree. Since the tree is self-balancing,
   this lookup will always be a O(lg n) operation.

   - parameter element: Element we want to find in the tree
   - returns: Element if it is found, nil otherwise.
   */
  func find(element: Element) -> Element? {
    if let node = self.findNode(element) {
      return node.value
    }
    return nil
  }
  /**
   Finds the node for an element in the AVL tree. Since the tree
   is self-balancing, this lookup will always be a O(lg n) operation.

   - parameter element: Element we want to find in the tree
   - returns: Node of the element if it is found, nil otherwise.
   */

  func findNode(element: Element) -> Node<Element>? {
    return findNode(element, node: _root)
  }

  private func findNode(element: Element, node: Node<Element>?) -> Node<Element>? {
    if let node = node {
      if node.value == element {
        return node
      } else if node.value > element {
        return findNode(element, node: node.leftNode)
      } else if node.value < element {
        return findNode(element, node: node.rightNode)
      }
    }
    return nil
  }
}

SequenceType extension:

extension BinarySearchTree: SequenceType {

  func generate() -> AnyGenerator<Element> {
    var _current: Node<Element>?
    return AnyGenerator {
      if _current != nil {
        _current = self.successor(_current!)
      } else {
        if self.root != nil {
          _current = self.minNode(self.root!)
        }
      }
      return _current != nil ? _current?.value : nil
    }
  }
}

CollectionType extension:

extension BinarySearchTree: CollectionType {

  var startIndex: Int {
    return 0
  }

  var endIndex: Int {
    return count
  }

  subscript(index: Int) -> Element {
    return self.findItemAtIndexUnsafe(index).value
  }

  private func findItemAtIndexUnsafe(index: Int) -> Node<Element> {
    var currentCount = index
    var result: Node<Element>?
    findItemAtIndex(&currentCount, node: root!, result: &result)
    return result!
  }

  func findItemAtIndex(inout index: Int, node: Node<Element>, inout result: Node<Element>?) {

    if node.leftNode != nil {
      findItemAtIndex(&index, node: node.leftNode!, result: &result)
    }

    if index == 0 {
      result = node
    }
    index -= 1

    if node.rightNode != nil {
      findItemAtIndex(&index, node: node.rightNode!, result: &result)
    }
  }

  func findItemAtIndex(index: Int) throws -> Node<Element> {
    if let root = root {
      var index = index
      var result: Node<Element>?
      findItemAtIndex(&index, node: root, result: &result)
      if let result = result {
        return result
      }
    }
    throw BinarySearchTreeError.OutOfBound
  }
}

ArrayLiteralConvertible extension:

extension BinarySearchTree: ArrayLiteralConvertible {
   convenience init(arrayLiteral: Element...) {
    self.init()
    for element in arrayLiteral {
      self.insert(element)
    }
  }
}

Sample test:

var tree = BinarySearchTree<Int>()

tree.insert(7)
tree.insert(3)
tree.insert(4)
tree.insert(9)
tree.insert(2)
tree.insert(1)

var node = tree.findNode(7)

for index in 0..<tree.count {
  print(tree[index])
}


var reverse = tree.reverse()

for value in reverse {
  print(value)
}

var newTree: BinarySearchTree = ["hello", "world", "this", "is", "a", "new", "day"]

_ = newTree.map {print($0)}

Output:

Hello, World!
1
2
3
4
7
9
9
7
4
3
2
1
a
day
hello
is
new
this
world
Program ended with exit code: 0

Answer 1

I have a major point of criticism regarding the "Swifty-ness" of your code: There is far too much "forced unwrapping" of optionals, and explicit comparing against nil. This can almost always be avoided by using

• Optional binding: if/while let myValue = myOptional { ... },
• optional chaining: myOptional?.property, or
• the nil-coalescing operator: myOptional ?? defaultValue.

See also When should I compare an optional value to nil? on Stack Overflow.

Example 1:

func minNode(root: Node<Element>) -> Node<Element> {
    var _current = root
    while _current.leftNode != nil {
        _current = _current.leftNode!
    }
    return _current
}

should be written with optional binding as

func minNode(root: Node<Element>) -> Node<Element> {
    var _current = root
    while let lNode = _current.leftNode {
        _current = lNode
    }
    return _current
}

It is even shorter because the leftNode property is accessed only once per loop iteration.

Example 2:

private func height (node: Node<Element>?) -> Int {
    return node != nil ? node!.height : -1
}

can be simplified with optional chaining and the nil-coalescing operator ??:

private func height (node: Node<Element>?) -> Int {
    return node?.height ?? -1
}

Example 3:

func insert(element: Element) {
    if let _ = _root {
        self.insert(element, currentNode: &_root!)
    } else {
        _root = Node(value: element)
        _count += 1
    }
}

Again forced unwrapping, and if let _ = _root is just if _root != nil in disguise. The above transformations do not work directly in this case because self.insert(element, currentNode: &_root!) modifies the last argument.

But actually the

private func insert(element: Element, inout currentNode: Node<Element>) -> Node<Element> {
    // ...
    return currentNode
}

does not really need an inout parameter because it returns the new root of the tree after the insertion. It can be rewritten as

private func insert(element: Element, currentNode: Node<Element>) -> Node<Element> {
    var currentNode = currentNode // make mutable copy
    // ... 
    return currentNode
}

And now the insert(element: Element) method can use optional binding as well:

func insert(element: Element) {
    if let node = _root {
        _root = self.insert(element, currentNode: node)
    } else {
        _root = Node(value: element)
        count += 1
    }
}

The same principle can be applied at many places in your code. I hope that the above examples can serve as a starting point to clean-up the handling of optionals.

---

The "public readonly, private read-write" properties "root" and "count" can be defined as

private(set) var root: Node<Element>?
private(set) var count: Int = 0 

which makes the additional accessor methods obsolete. Actually root need not be publicly visible at all.

The value property of Node should be declared as a constant because modifying the value makes no sense after a node as been inserted in the tree.

My suggestion would be to fix the "optional handling" first and then review the code in a second iteration.

Answer 2

To the very good analysis by Martin I would add that I feel that variable shadowing, which seems to belong to the Swift style (at least in the let statements) is in fact unfortunate. The example:

func findItemAtIndex(index: Int) throws -> Node<Element> {
    if let root = root {
      var index = index
...

Instance variable root and the parameter index are being shadowed by local variables. In this case it doesn't do much harm because the shadowing happens directly at let and almost immmediately after index parameter's definition and the method is short. In fact the method contains one more shadowing:

if let result = result {
    return result
}

which is causing me to scratch my head at first glance.

In longer and more complicated methods I would argue against shadowing, perhaps to be consistent against them totally. I know it is a matter of personal style and that there are accepted style guides and I know of the advantages. My arguments is that some of the most important aspects of producing code is readability and understandability - above all for all that come after us to change our code. And variable shadowing harms both aspects.

It helps reduce name space pollution by unnecessary names, the difference between the optional and unwrapped value (which is essential in my eyes) vanishes though. And not seeing the location where shadowing happens brings me to wonder why is this thing being accessed in that way and what is really being changed by changing this or that variable.

Another one: the balance method is unused, is declared private so can be removed.