Binary search tree iterator in Swift


  • 0
    S

    Binary search tree defined using enum and its iterator by extending IteratorProtocol

    enum BinarySearchTree <T: Comparable> {
        case empty
        indirect case node (BinarySearchTree<T>, T, BinarySearchTree<T>)
        
        var count: Int {
            switch self {
            case let .node(left, _, right):
                return left.count + 1 + right.count
            case .empty:
                return 0
            }
        }
        
        private func newTreeWithInsertedValue(newValue: T) -> BinarySearchTree {
            switch self {
            case .empty:
                return .node(.empty, newValue, .empty)
            case let .node(left, value, right):
                if newValue < value {
                    return .node(left.newTreeWithInsertedValue(newValue: newValue), value, right)
                } else {
                    return .node(left, value, right.newTreeWithInsertedValue(newValue: newValue))
                }
            }
        }
        
        mutating func insert(newValue: T) {
            self = newTreeWithInsertedValue(newValue: newValue)
        }
        
        func contains(searchValue: T) -> Bool {
            switch self {
            case .empty:
                return false
            case let .node(left, value, right):
                if searchValue == value { return true }
                if searchValue < value {
                    return left.contains(searchValue: searchValue)
                } else {
                    return right.contains(searchValue: searchValue)
                }
            }
        }
        
        func inorderTraversal(process: (T) -> ()) {
            switch self {
            case .empty:
                return
            case let .node(left, value, right):
                left.inorderTraversal(process: process)
                process(value)
                right.inorderTraversal(process: process)
            }
        }
    }
    
    extension BinarySearchTree: Sequence {
        func makeIterator() -> BinarySearchTreeIterator<T> {
            return BinarySearchTreeIterator(self)
        }
        
        struct BinarySearchTreeIterator<T: Comparable>: IteratorProtocol {
            var tree: BinarySearchTree<T>
            var stack = [(T, BinarySearchTree<T>)]()
            
            init(_ tree: BinarySearchTree<T>) {
                self.tree = tree
            }
            
            mutating func next() -> T? {
                while case let .node(left, value, right) = tree {
                    stack.append((value, right))
                    tree = left
                }
                
                guard let (element, tree) = stack.popLast() else { return nil }
                self.tree = tree
                return element
            }
        }
    }
    
    
    
    // var anotherTree: BinarySearchTree<Int> = BinarySearchTree.empty
    // [5,2,4,8,3,1,9,11,20,10].forEach { anotherTree.insert(newValue: $0) }
    
    // for element in anotherTree {
        // print(element)
    // }
    
    // anotherTree.reduce(0, +)
    
    
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.