Skip to content

Commit 2f96bbb

Browse files
jsbeanjsbean
authored
Add AVLTree (#166)
1 parent f52308c commit 2f96bbb

File tree

4 files changed

+284
-2
lines changed

4 files changed

+284
-2
lines changed

‎Sources/DataStructures/AVLTree.swift

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
//
2+
// AVLTree.swift
3+
// DataStructures
4+
//
5+
// Created by James Bean on 9/9/18.
6+
//
7+
8+
import Destructure
9+
10+
/// Copy-on-write self-balancing binary search tree.
11+
public struct AVLTree <Key: Comparable, Value> {
12+
13+
// MARK: - Instance Properties
14+
15+
@usableFromInline
16+
var root: Node?
17+
18+
/// Creates an `AVLTree` with the given `root` node.
19+
@usableFromInline
20+
init(root: Node? = nil) {
21+
self.root = root
22+
}
23+
}
24+
25+
extension AVLTree {
26+
27+
// MARK: - Nested Types
28+
29+
/// Node for an `AVLTree` which contains a `key`, and `value`, along with the `left` and `right`
30+
/// child nodes (if they exist), as well as the its `height`.
31+
public final class Node {
32+
33+
// MARK: - Instance Properties
34+
35+
// MARK: Key and Value
36+
37+
/// The key used to keep this node sorted within an `AVLTree`.
38+
public var key: Key
39+
40+
/// The payload contained herein.
41+
public var value: Value
42+
43+
// MARK: Child nodes
44+
45+
/// The node storing values with keys less than this one.
46+
public var left: Node?
47+
48+
/// The node storing values with keys greater than this one.
49+
public var right: Node?
50+
51+
// MARK: Position
52+
53+
/// The height of this node within an `AVLTree`.
54+
public var height: Int
55+
56+
// MARK: - Initializers
57+
58+
/// Creates an `AVLTree.Node` with the given `key`, `value`, and `height`.
59+
@inlinable
60+
public init(key: Key, value: Value, height: Int = 1) {
61+
self.key = key
62+
self.value = value
63+
self.height = height
64+
}
65+
}
66+
}
67+
68+
extension AVLTree {
69+
70+
// MARK: - Initializers
71+
72+
/// Creates an `AVLTree` with the given `key` and `value` for the `root` node.
73+
@inlinable
74+
public init(key: Key, value: Value) {
75+
self.root = Node(key: key, value: value)
76+
}
77+
78+
/// Creates an `AVLTree` with the given `sequence` of key-value pairs.
79+
@inlinable
80+
public init <S> (_ sequence: S) where S: Sequence, S.Element == (Key,Value) {
81+
self.init()
82+
sequence.forEach { key,value in insert(value, forKey: key) }
83+
}
84+
}
85+
86+
extension AVLTree {
87+
88+
// MARK: - Computed Properties
89+
90+
/// - Returns: An array of key-value pairs in sorted order.
91+
@inlinable
92+
public var inOrder: [(Key,Value)] {
93+
func traverse(_ node: Node, into result: inout [(Key,Value)]) -> [(Key,Value)] {
94+
let left = node.left.map { traverse($0, into: &result ) } ?? []
95+
let right = node.right.map { traverse($0, into: &result) } ?? []
96+
return left + [(node.key,node.value)] + right
97+
}
98+
var result: [(Key,Value)] = []
99+
guard let root = root else { return result }
100+
return traverse(root, into: &result)
101+
}
102+
}
103+
104+
extension AVLTree {
105+
106+
// MARK: - Instance Methods
107+
108+
/// Inserts the given `value` for the given `key`.
109+
@inlinable
110+
public mutating func insert(_ value: Value, forKey key: Key) {
111+
ensureUnique()
112+
if let node = root {
113+
root = AVLTree.insert(value, forKey: key, into: node)
114+
} else {
115+
root = Node(key: key, value: value)
116+
}
117+
}
118+
119+
@usableFromInline
120+
mutating func ensureUnique() {
121+
if !isKnownUniquelyReferenced(&root) {
122+
self = AVLTree(root: root)
123+
}
124+
}
125+
126+
// FIXME: If Value is RangeRepleceableCollection (or Additive?), instead of returning `node`
127+
// upon interval match, concatenate incoming `value` with extant value.
128+
@usableFromInline
129+
@discardableResult
130+
static func insert(_ value: Value, forKey key: Key, into node: Node? = nil) -> Node? {
131+
guard let node = node else { return Node(key: key, value: value) }
132+
if node.key > key {
133+
node.left = insert(value, forKey: key, into: node.left)
134+
} else if node.key < key {
135+
node.right = insert(value, forKey: key, into: node.right)
136+
} else {
137+
return node
138+
}
139+
return balance(node, with: key)
140+
}
141+
142+
static func balance(_ node: Node, with key: Key) -> Node {
143+
node.updateHeight()
144+
let balance = balanceFactor(node.left, node.right)
145+
if balance > 1 && key < node.left!.key {
146+
return rotateRight(node)
147+
} else if balance < -1 && key > node.right!.key {
148+
return rotateLeft(node)
149+
} else if balance > 1 && key > node.left!.key {
150+
node.left = rotateLeft(node.left!)
151+
return rotateRight(node)
152+
} else if balance < -1 && key < node.right!.key {
153+
node.right = rotateRight(node.right!)
154+
return rotateLeft(node)
155+
}
156+
return node
157+
}
158+
159+
// TODO: removeValue(forKey key: Key) -> Node?
160+
161+
static func balanceFactor(_ left: Node?, _ right: Node?) -> Int {
162+
return (left?.height ?? -1) - (right?.height ?? -1)
163+
}
164+
165+
static func rotateLeft(_ node: Node) -> Node {
166+
guard let newRoot = node.right else { return node }
167+
node.right = newRoot.left
168+
newRoot.left = node
169+
node.updateHeight()
170+
newRoot.updateHeight()
171+
return newRoot
172+
}
173+
174+
static func rotateRight(_ node: Node) -> Node {
175+
guard let newRoot = node.left else { return node }
176+
node.left = newRoot.right
177+
newRoot.right = node
178+
node.updateHeight()
179+
newRoot.updateHeight()
180+
return newRoot
181+
}
182+
}
183+
184+
extension AVLTree.Node {
185+
186+
func updateHeight() {
187+
self.height = max(left?.height ?? 0, right?.height ?? 0) + 1
188+
}
189+
}
190+
191+
extension AVLTree.Node: Equatable where Value: Equatable {
192+
193+
// MARK: - Equatable
194+
195+
/// - Returns: `true` if the two given `AVLTree.Node` values are equivalent. Otherwise, `nil`.
196+
public static func == (lhs: AVLTree.Node, rhs: AVLTree.Node) -> Bool {
197+
guard lhs.key == rhs.key, lhs.value == rhs.value else { return false }
198+
return lhs.left == rhs.left && lhs.right == rhs.right
199+
}
200+
}
201+
202+
extension AVLTree: Equatable where Value: Equatable { }
203+
204+
extension AVLTree: ExpressibleByArrayLiteral {
205+
206+
// MARK: - ExpressibleByArrayLiteral
207+
208+
/// Creates an `AVLTree` with an array literal of key-value pairs.
209+
public init(arrayLiteral elements: (Key,Value)...) {
210+
self.init(elements)
211+
}
212+
}

‎Tests/DataStructuresTests/AVLTreeTests.swift

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
//
2+
// AVLTreeTests.swift
3+
// DataStructuresTests
4+
//
5+
// Created by James Bean on 9/9/18.
6+
//
7+
8+
import XCTest
9+
@testable import DataStructures
10+
11+
class AVLTreeTests: XCTestCase {
12+
13+
func testInitSequence() {
14+
let _: AVLTree = [(3,"d"),(1,"b"),(2,"c"),(0,"a")]
15+
}
16+
17+
func testSingleNodeHeight() {
18+
let tree: AVLTree = [(0,"0")]
19+
let expected = AVLTree(root: AVLTree.Node(key: 0, value: "0"))
20+
XCTAssertEqual(tree, expected)
21+
}
22+
23+
func testTwoNodes() {
24+
let tree: AVLTree = [(0,"0"),(1,"1")]
25+
let child = AVLTree.Node(key: 1, value: "1", height: 1)
26+
let parent = AVLTree.Node(key: 0, value: "0", height: 2)
27+
parent.right = child
28+
let expected = AVLTree(root: parent)
29+
XCTAssertEqual(tree, expected)
30+
}
31+
32+
func testThreeNodes() {
33+
let tree: AVLTree = [(0,"0"),(1,"1"),(2,"2")]
34+
let left = AVLTree.Node(key: 0, value: "0", height: 2)
35+
let parent = AVLTree.Node(key: 1, value: "1", height: 1)
36+
let right = AVLTree.Node(key: 2, value: "2", height: 1)
37+
parent.left = left
38+
parent.right = right
39+
let expected = AVLTree(root: parent)
40+
XCTAssertEqual(tree, expected)
41+
}
42+
43+
func testManyValuesIncreasing() {
44+
let values = (0..<100_000).map { ($0,$0) }
45+
let _ = AVLTree(values)
46+
}
47+
48+
func testManyValuesDecreasing() {
49+
let values = (0..<100_000).reversed().map { ($0,$0) }
50+
let _ = AVLTree(values)
51+
}
52+
53+
func testManyValuesRandom() {
54+
let values = (0..<1_000_000).map { ($0,$0) }.shuffled()
55+
let _ = AVLTree(values)
56+
}
57+
}

‎Tests/DataStructuresTests/XCTestManifests.swift

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,17 @@
11
import XCTest
22

3+
extension AVLTreeTests {
4+
static let __allTests = [
5+
("testInitSequence", testInitSequence),
6+
("testManyValuesDecreasing", testManyValuesDecreasing),
7+
("testManyValuesIncreasing", testManyValuesIncreasing),
8+
("testManyValuesRandom", testManyValuesRandom),
9+
("testSingleNodeHeight", testSingleNodeHeight),
10+
("testThreeNodes", testThreeNodes),
11+
("testTwoNodes", testTwoNodes),
12+
]
13+
}
14+
315
extension ArrayExtensionsTests {
416
static let __allTests = [
517
("testInsertingAtIndexAtBeginning", testInsertingAtIndexAtBeginning),
@@ -464,6 +476,7 @@ extension ZipToLongestTests {
464476
#if !os(macOS)
465477
public func __allTests() -> [XCTestCaseEntry] {
466478
return [
479+
testCase(AVLTreeTests.__allTests),
467480
testCase(ArrayExtensionsTests.__allTests),
468481
testCase(BimapTests.__allTests),
469482
testCase(BinaryHeapTests.__allTests),

‎Tests/LinuxMain.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,17 @@
11
import XCTest
22

3+
import DataStructuresTests
34
import AlgebraPerformanceTests
45
import AlgebraTests
5-
import DataStructuresTests
66
import DataStructuresPerformanceTests
77
import AlgorithmsTests
88
import DestructureTests
99
import AlgorithmsPerformanceTests
1010

1111
var tests = [XCTestCaseEntry]()
12+
tests += DataStructuresTests.__allTests()
1213
tests += AlgebraPerformanceTests.__allTests()
1314
tests += AlgebraTests.__allTests()
14-
tests += DataStructuresTests.__allTests()
1515
tests += DataStructuresPerformanceTests.__allTests()
1616
tests += AlgorithmsTests.__allTests()
1717
tests += DestructureTests.__allTests()

0 commit comments

Comments
 (0)