Tree.swift

import UIKit

typealias TreeNode = TreeElement<Int>

enum Color { case r, b }

indirect enum TreeElement<T: Comparable> {
    
    case empty
    case node(Color, TreeElement<T>, T, TreeElement<T>)
    
    public var count: Int {
        switch self {
        case let .node(_, left, _ , right):
            return left.count + 1 + right.count
        case .empty:
            return 0
        }
    }
    
    init() {
        self = .empty
    }
    
    init(_ x: T, color: Color = .b, left: TreeElement<T> = .empty, right: TreeElement<T> = .empty) {
        self = .node(color, left, x, right)
    }
    
    private func balance<T>(_ tree: TreeElement<T>) -> TreeElement<T> {
        switch tree {
        case let .node(.b, .node(.r, .node(.r, a, x, b), y, c), z, d):
            return .node(.r, .node(.b, a, x, b),y,.node(.b, c, z, d))
        case let .node(.b, .node(.r, a, x, .node(.r, b, y, c)), z, d):
            return .node(.r, .node(.b, a, x, b),y,.node(.b, c, z, d))
        case let .node(.b, a, x, .node(.r, .node(.r, b, y, c), z, d)):
            return .node(.r, .node(.b, a, x, b),y,.node(.b, c, z, d))
        case let .node(.b, a, x, .node(.r, b, y, .node(.r, c, z, d))):
            return .node(.r, .node(.b, a, x, b),y,.node(.b, c, z, d))
        default:
            return tree
        }
    }
    
    func insert<T>(_ into: TreeElement<T>, _ x: T) -> TreeElement<T> {
        guard case let .node(c, l, y, r) = into else { return TreeElement<T>(x, color: .r) }
        if x < y { return balance(TreeElement<T>(y, color: c, left: insert(l,x), right: r)) }
        if y < x { return balance(TreeElement<T>(y, color: c, left: l, right: insert(r, x))) }
        return into
    }
}

var first = TreeNode(0, color: .r, left:  TreeNode.node(.b, TreeNode.empty, 10, .empty), right: .empty)
var third = TreeNode.node(.b, TreeNode.empty, 15, TreeNode.empty)
var second = TreeNode(1, color: .b, left: third, right: first)
let tree = TreeNode(2, color: .b, left: second, right: first)

tree.insert(second, 25)
tree.insert(TreeNode(12), 5)
print(tree.count)