Swift遞迴列舉與紅黑樹

bestswifter發表於2017-12-14
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以示鼓勵。

相關文章