Red-Black Tree implementation in Swift

RedBlackTree.swift

//
//  main.swift
//  Beachline
//
//  Created by Oleksandr Glagoliev on 1/4/20.
//  Copyright © 2020 Oleksandr Glagoliev. All rights reserved.
//

import Foundation

/// A red-black tree is a binary tree that satisfies the following red-black properties:
/// 1. Every node is either red or black.
/// 2. The root is black.
/// 3. Every leaf (NIL) is black.
/// 4. If a node is red, then both its children are black.
/// 5. For each node, all simple paths from the node to descendant leaves contain the same number of black nodes.

final class RedBlackNode<K: Comparable> {
    var isBlack: Bool = true
    var key: K? = nil
    
    var right: RedBlackNode<K>?
    var left: RedBlackNode<K>?
    unowned var parent: RedBlackNode<K>?
}

final class RedBlackTree<K: Comparable> {
    private var root: RedBlackNode<K>?
    private let sentinel: RedBlackNode<K>
    
    init() {
        sentinel = RedBlackNode<K>()
        root = sentinel
        root?.left = sentinel
        root?.right = sentinel
    }
    
    private func minimum(_ x: RedBlackNode<K>) -> RedBlackNode<K> {
        var x = x
        while x.left !== sentinel {
            x = x.left!
        }
        return x
    }
    
    private func maximum(_ x: RedBlackNode<K>) -> RedBlackNode<K> {
        var x = x
        while x.right !== sentinel {
            x = x.right!
        }
        return x
    }
    
    func search(x: RedBlackNode<K>, k: K) -> RedBlackNode<K> {
        var x = x
        while x !== sentinel && k != x.key {
            if k < x.key! {
                x = x.left!
            } else {
                x = x.right!
            }
        }
        return x
    }
    
    func successor(x: RedBlackNode<K>) -> RedBlackNode<K>? {
        var x = x
        if x.right !== sentinel {
            return minimum(x.right!)
        }
        var successor = x.parent
        while successor !== sentinel && x === successor!.right {
            x = successor!
            successor = successor?.parent
        }
        return successor
    }
    
    func predecessor(_ x: RedBlackNode<K>) -> RedBlackNode<K>? {
        var x = x
        if x.left !== sentinel {
            return minimum(x.left!)
        }
        var predecessor = x.parent
        while predecessor !== sentinel && x === predecessor!.left {
            x = predecessor!
            predecessor = predecessor?.parent
        }
        return predecessor
    }
    
    private func insert(_ z: RedBlackNode<K>) {
        var y = sentinel
        var x = root
        while x !== sentinel {
            y = x!
            if z.key! < x!.key! {
                x = x!.left
            } else {
                x = x!.right
            }
        }
        z.parent = y
        if y === sentinel {
            root = z
        } else if z.key! < y.key! {
            y.left = z
        } else {
            y.right = z
        }
        z.left = sentinel
        z.right = sentinel
        z.isBlack = false
        insertFixup(z)
    }
    
    private func delete(_ z: RedBlackNode<K>) {
        var y = z
        var yOriginalColorIsBlack = y.isBlack
        var x: RedBlackNode<K>!
        if z.left === sentinel {
            x = z.right
            transplant(z, z.right!)
        } else if z.right === sentinel {
            x = z.left
            transplant(z, z.left!)
        } else {
            y = minimum(z.right!)
            yOriginalColorIsBlack = y.isBlack
            x = y.right
            if y.parent === z {
                x.parent = y
            } else {
                transplant(y, y.right!)
                y.right = z.right
                y.right!.parent = y
            }
            transplant(z, y)
            y.left = z.left
            y.left?.parent = y
            y.isBlack = z.isBlack
        }
        if yOriginalColorIsBlack {
            deletionFixup(x)
        }
    }
    
    private func insertFixup(_ z: RedBlackNode<K>) {
        var z = z
        while z.parent!.isBlack == false {
            if z.parent === z.parent!.parent!.left {
                let y = z.parent!.parent!.right
                if y?.isBlack == false {
                    z.parent?.isBlack = true
                    y?.isBlack = true
                    z.parent?.parent?.isBlack = false
                    z = z.parent!.parent!
                } else {
                    if z === z.parent!.right {
                        z = z.parent!
                        leftRotate(z)
                    }
                    z.parent!.isBlack = true
                    z.parent!.parent!.isBlack = false
                    rightRotate(z.parent!.parent!)
                }
            } else {
                let y = z.parent!.parent!.left
                if y?.isBlack == false {
                    z.parent?.isBlack = true
                    y?.isBlack = true
                    z.parent!.parent!.isBlack = false
                    z = z.parent!.parent!
                } else {
                    if z === z.parent!.left {
                        z = z.parent!
                        rightRotate(z)
                    }
                    z.parent?.isBlack = true
                    z.parent!.parent!.isBlack = false
                    leftRotate(z.parent!.parent!)
                }
            }
        }
        root?.isBlack = true
    }
    
    private func deletionFixup(_ x: RedBlackNode<K>) {
        var x = x
        while x !== root && x.isBlack == true {
            if x === x.parent?.left {
                var w = x.parent?.right
                if w?.isBlack == false {
                    w?.isBlack = true
                    x.parent?.isBlack = false
                    leftRotate(x.parent!)
                    w = x.parent?.right
                }
                if w?.left?.isBlack == true && w?.right?.isBlack == true {
                    w?.isBlack = false
                    x = x.parent!
                } else {
                    if w?.right?.isBlack == true {
                        w?.left?.isBlack = true
                        w?.isBlack = false
                        rightRotate(w!)
                        w = x.parent?.right
                    }
                    w?.isBlack = x.parent!.isBlack
                    x.parent?.isBlack = true
                    w?.right?.isBlack = true
                    leftRotate(x.parent!)
                    x = root!
                }
            } else {
                var w = x.parent?.left
                if w?.isBlack == false {
                    w?.isBlack = true
                    x.parent?.isBlack = false
                    rightRotate(x.parent!)
                    w = x.parent?.left
                }
                if w?.right?.isBlack == true && w?.left?.isBlack == true {
                    w?.isBlack = false
                    x = x.parent!
                } else {
                    if w?.left?.isBlack == true {
                        w?.right?.isBlack = true
                        w?.isBlack = true
                        leftRotate(w!)
                        w = x.parent?.left
                    }
                    w?.isBlack = x.parent!.isBlack
                    x.parent?.isBlack = true
                    w?.left?.isBlack = true
                    rightRotate(x.parent!)
                    x = root!
                }
            }
        }
        x.isBlack = false
    }
    
    private func rightRotate(_ x: RedBlackNode<K>) {
        let y = x.left
        x.left = y?.right //reassign g
        if y?.right !== sentinel {
            y?.right?.parent = x //reassign g's parent
        }
        y?.parent = x.parent //reassign the subtree's parent
        if x.parent === sentinel { //if the root is the root of the tree
            root = y
        } else if x === x.parent?.right { //reassign the original root parent's child node
            x.parent?.right = y
        } else {
            x.parent?.left = y
        }
        //establish parent/child relationship between old and new root nodes
        y?.right = x
        x.parent = y
    }
    
    private func leftRotate(_ x: RedBlackNode<K>) {
        let y = x.right
        x.right = y?.left
        if y?.left !== sentinel {
            y?.left?.parent = x
        }
        y?.parent = x.parent
        if x.parent === sentinel {
            root = y
        } else if x === x.parent?.left {
            x.parent?.left = y
        } else {
            x.parent?.right = y
        }
        y?.left = x
        x.parent = y
    }
    
    private func transplant(_ u: RedBlackNode<K>, _ v: RedBlackNode<K>) {
        if u.parent === sentinel {
            root = v
        } else if u === u.parent?.left {
            u.parent?.left = v
        } else {
            u.parent?.right = v
        }
        v.parent = u.parent
    }
}

extension RedBlackTree {
    public func insertKey(key: K) {
        let newNode = RedBlackNode<K>()
        newNode.key = key
        insert(newNode)
    }
    
    public func deleteKey(key: K) {
        let nodeToDelete = search(x: root!, k: key)
        if nodeToDelete !== sentinel {
            delete(nodeToDelete)
        }
    }
    
    public func findKey(key: K) -> RedBlackNode<K>? {
        let foundNode = search(x: root!, k: key)
        return foundNode === sentinel
            ? nil
            : foundNode
    }
    
    public var min: RedBlackNode<K>? {
        guard let r = root, r !== sentinel else {
            return nil
        }
        return minimum(r)
    }
    
    public var max: RedBlackNode<K>? {
        guard let r = root, r !== sentinel else {
            return nil
        }
        return maximum(r)
    }
}

extension RedBlackTree {
    func printDOTFormat() {
        print("digraph G {")
        if let root = root {
            toDOT(node: root)
        }
        print("}")
    }
    
    private func toDOT(node: RedBlackNode<K>) {
        guard node !== sentinel else { return }
        let curr = toString(node)
        if let l = node.left {
            let left = toString(l)
            print("\"\(curr)\"->\"\(left)\"[taillabel = \"L\"]")
            toDOT(node: l)
        }
        if let r = node.right {
            let right = toString(r)
            print("\"\(curr)\"->\"\(right)\"[taillabel = \"R\"]")
            toDOT(node: r)
        }
    }
    
    private func toString(_ node: RedBlackNode<K>) -> String {
        if let key = node.key {
            return "\(key)"
        } else {
            return "nil"
        }
    }
}

let rbTree = RedBlackTree<Int>()

//rbTree.insertKey(key: 1)
//rbTree.insertKey(key: 2)
//rbTree.insertKey(key: 3)
//rbTree.insertKey(key: 4)
//rbTree.insertKey(key: 5)
//
//rbTree.deleteKey(key: 4)

//print(rbTree.findKey(key: 3))
//print(rbTree.findKey(key: 6))
//print(rbTree.min?.key)
//print(rbTree.max?.key)


let numInsertions = 30
var numbers: [Int] = []
(0..<numInsertions).forEach {
//    let key: Int = Int.random(in: 0...1000000)
    numbers.append($0)
}

let start = CFAbsoluteTimeGetCurrent()
numbers.forEach {
    rbTree.insertKey(key: $0)
//    print(rbTree.min.k, rbTree.max?.key)
//    print(rbTree.min?.key)
//    print(rbTree.max?.key)
    
}
rbTree.printDOTFormat()
var i = 0
numbers.forEach {
    if $0 % 2 == 0 {
        rbTree.deleteKey(key: $0)
    }
    i+=1
}
//rbTree.printDOTFormat()
let diff = CFAbsoluteTimeGetCurrent() - start
print("It Took \(diff) seconds to insert \(numInsertions)")