Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Red-Black Tree implementation in 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)")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment