Swift遞歸枚舉與紅黑樹

爲了節省篇幅，本文省略了紅黑樹部分大量的註釋，只展現代碼部分，完整的項目代碼在個人Github上：嵌套枚舉與紅黑樹，若是以爲不錯還但願隨手點一個star。

在使用OC時，枚舉是一種很是簡單易用的數據結構，它能夠經過某些自定義，具備顯著語義的符號來定義某些狀態，使代碼更加清晰易懂：

typedef NS_ENUM(NSInteger, MyEnum) {
ValueA = 1,  // 若是不寫值，默認爲1，這是枚舉的原始值
ValueB, // 默認是ValueA的值加1
ValueC, // 默認是ValueB的值加1，以此類推
ValueD
};
複製代碼

關聯值

在Swift，枚舉有了更多的用法。除了能夠具備原始值外，枚舉還能夠具備關聯值(Associated Values)。關聯值容許枚舉成員包含更多自定義信息，知足咱們的使用。好比有一個枚舉類型能夠接受數字密碼或字符串密碼：

enum Password {
case DeigtPassword(Int)
case StringPassword(String)
}
複製代碼

在枚舉成員的名字後面多了一個括號，這個括號裏的類型就是枚舉成員關聯值的類型。也就是說，枚舉成員DeigtPassword具備一個Int類型的關聯值，枚舉成員StringPassword具備一個String類型的關聯值。枚舉值的建立以下：

var password = Password.DeigtPassword(123456)
password = .StringPassword("123456")
複製代碼

同一時刻，一個枚舉變量，也就是這裏的password，只能存儲一個枚舉成員和它的關聯值。讀取關聯值的語法以下：

switch password {
case .DeigtPassword(let digit): print("這是數字密碼： \(digit)")
case .StringPassword(let digit): print("這是字符串密碼： \(digit)")
}
複製代碼

這說明關聯值能夠經過case let語句取出。

遞歸枚舉

在咱們的印象中，有序二叉樹這樣的數據結構顯然是應該用類來定義的：

class Tree<Element: Comparable> {
let value: Element?
// entries < value go on the left
let left: Tree<Element>?
// entries > value go on the right
let right: Tree<Element>?
}
複製代碼

可是，因爲根節點、左子樹和右子樹都有可能爲空，因此類中的每個成員都是可選類型，這致使代碼看上去很是醜，也不方便實現。使用枚舉能夠很容易的解決這個問題：

enum Tree<Element: Comparable> {
case Empty
case Node(Tree<Element>,Element,Tree<Element>)
}
複製代碼

這時候，枚舉成員Empty表示樹爲空，若是樹不爲空，那麼他就具備根節點和左右子樹(若是爲空，那麼值就是Tree.Empty)。這裏的枚舉成員Node就具備關聯值，不過它關聯的不是一個變量，而是三個變量。

在枚舉成員Node的關聯值中，有兩個值的類型是Tree<Element>，這種枚舉的嵌套是不容許的。其實簡單回憶一下，類成員的嵌套也是不容許的：

class A {
var a = A()
}

var a = A()  // 報錯
複製代碼

好在咱們能夠爲枚舉類型添加indirect關鍵字來實現這種嵌套關係：

indirect enum Tree<Element: Comparable> {
case Empty
case Node(Tree<Element>,Element,Tree<Element>)
}
複製代碼

紅黑樹

這樣作還有一個潛在的問題，由於二叉樹是有序的，當咱們插入一組有序的數組時，二叉樹變得很是不平衡。好比要插入的數組是[1, 2, 3, 4]：

• 插入元素1： 根節點爲1
• 插入元素2： 根節點爲1，根節點的右子樹爲2
• 插入元素3： 根節點爲1，根節點的右子樹爲2，2的右子樹爲3
• 插入元素4： 根節點爲1，根節點的右子樹爲2，2的右子樹爲3，3的右子樹爲4
• ……

這樣的二叉樹，進行二叉搜索時的時間複雜度爲O(n)，這顯然不是咱們但願的結果。因此能夠把它用紅黑樹實現：

enum Color { case R, B }

indirect enum Tree<Element: Comparable> {
case Empty
case Node(Color,Tree<Element>,Element,Tree<Element>)

init() { self = .Empty }

init(_ x: Element, color: Color = .B,
left: Tree<Element> = .Empty, right: Tree<Element> = .Empty)
{
self = .Node(color, left, x, right)
}
}
複製代碼

Contains方法

接下來咱們實現一個contains方法，用來判斷某個元素是否存在於樹中：

extension Tree {
func contains(x: Element) -> Bool {
guard case let .Node(_,left,y,right) = self
else { return false }

if x < y { return left.contains(x) }
if y < x { return right.contains(x) }
return true
}
}
複製代碼

guard語句的使用很巧妙，它判斷這棵樹當前是否是空，若是爲空則返回false，不然經過case let讀取出枚舉變量中的關聯值。

Insert方法

接下來實現的是insert方法，用於向樹中插入一個新的節點：

extension Tree {
func insert(x: Element) -> Tree {
guard case let .Node(_,l,y,r) = ins(self, x)
else { fatalError("ins should never return an empty tree") }

return .Node(.B,l,y,r)
}
}
複製代碼

代碼的主要邏輯在ins方法中實現：

private func ins<T>(into: Tree<T>, _ x: T) -> Tree<T> {
guard case let .Node(c, l, y, r) = into
else { return Tree(x, color: .R) }

if x < y { return balance(Tree(y, color: c, left: ins(l,x), right: r)) }
if y < x { return balance(Tree(y, color: c, left: l, right: ins(r, x))) }
return into
}
複製代碼

這裏面還有一個balance方法，主要處理節點位置和顏色的變化，具體代碼實現請參考個人Github：嵌套枚舉與紅黑樹。

中序遍歷

爲了實現樹的中序遍歷，首先要實現SequenceType協議，這樣就能夠用for in語法遍歷樹了。關於SequenceType協議，您能夠參考個人另外一篇文章：深刻探究Swift數組背後的協議、方法、拓展，代碼實現以下：

extension Tree: SequenceType {
func generate() -> AnyGenerator<Element> {
var stack: [Tree] = []
var current: Tree = self

return anyGenerator { _ -> Element? in
while true {
// if there's a left-hand node, head down it
if case let .Node(_,l,_,_) = current {
stack.append(current)
current = l
}
// if there isn’t, head back up, going right as
// soon as you can:
else if !stack.isEmpty, case let .Node(_,_,x,r) = stack.removeLast() {
current = r
return x
}
else {
// otherwise, we’re done
return nil
}
}
}
}
}
複製代碼

數組字面量

最後一個功能是根據數組字面量建立樹。這個功能很是容易實現，由於咱們此前已經實現了insert方法，如今只要再實現一下ArrayLiteralConvertible協議便可：

extension Tree: ArrayLiteralConvertible {
init <S: SequenceType where S.Generator.Element == Element>(_ source: S) {
self = source.reduce(Tree()) { $0.insert($1) }
}

init(arrayLiteral elements: Element...) {
self = Tree(elements)
}
}
複製代碼

在示範項目中，我對這個功能進行了一些測試：

let alphabet = Tree("the quick brown fox jumps over the lazy dog".characters)
for node in alphabet {
print(node)  // 打印從a到z的全部字母
}
複製代碼

Where To Go

紅黑樹的實現遠沒有結束，不過這篇博客的重點只是但願經過具體的例子講解嵌套關聯類型的使用。因此並無涉及紅黑樹的刪除等操做。若是您感興趣，歡迎提交代碼或issue到Github，若是以爲代碼對您有用過，也歡迎隨手點一個Star以示鼓勵。