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以示鼓励。