Trie.swift

extension Trie {
    class Node {
        var element: Element?
        var children: [Path: Node]

        init() {
            self.element = nil
            self.children = [:]
        }

        /// 再帰的に直列なNode列を作る(string[index + 1]を根とする部分木を作成)
        fileprivate init<C: Collection>(_ element: Element, path: C, after index: C.Index) where C.Element == Path {
            let nextIndex = path.index(after: index)
            if nextIndex == path.endIndex {
                self.element = element
                self.children = [:]
            } else {
                self.element = nil
                self.children = [path[nextIndex]: Node(element, path: path, after: nextIndex)]
            }
        }

        /// そのノードの下についてる文字列一覧を返す
        func subLeaves() -> [Element] {
            guard !children.isEmpty else {
                return element.map { [$0] } ?? []
            }
            return children.values.flatMap { child -> [Element] in
                return child.subLeaves()
            }
        }

        /// 自身を根とする部分木のから、対応するprefix分動いたノードを探す
        func findNode<C: Collection>(prefix: C) -> Node? where C.Element == Path {
            guard let firstCharacter = prefix.first else {
                return self
            }
            let remaining = prefix.dropFirst()
            if let node = children[firstCharacter]?.findNode(prefix: remaining) {
                return node
            }
            return nil
        }

        /// 改行を少なく表示
        fileprivate func showPretty(indent: Int) {
            guard !children.isEmpty else {
                print("") // 葉で改行
                return
            }
            var isFirst: Bool = true
            for node in children {
                if isFirst {
                    print("\(node.key)", terminator: "")
                    isFirst = false
                } else {
                    let spaces = [String](repeating: " ", count: indent).joined()
                    print(spaces + "\(node.key)", terminator: "")
                }
                node.value.showPretty(indent: indent + 1)
            }
        }
    }
}

final class Trie<Element, Path: Hashable> {
    var rootNode: Node

    init() {
        rootNode = Node()
    }

    func insert<C: Collection>(_ element: Element, at path: C) where C.Element == Path {
        var currentNode = rootNode
        for index in path.indices {
            let character = path[index]
            if let node = currentNode.children[character] {
                currentNode = node
                continue
            } else {
                currentNode.children[character] = Node(element, path: path, after: index)
                break
            }
        }
    }

    /// prefixで検索
    func find<C: Collection>(prefix: C) -> [Element] where C.Element == Path {
        guard let node = rootNode.findNode(prefix: prefix) else {
            return []
        }
        return node.subLeaves()
    }

    /// 後ろ何文字かが一致するものを検索（時間かかるので文字数減らしてから渡すこと）
    func find<C: Collection>(from string: C) -> [Element] where C.Element == Path {
        return string.indices.flatMap { index -> [Element] in
            guard let node = rootNode.findNode(prefix: string[index...]) else {
                return []
            }
            return node.subLeaves()
        }
    }

    func showPretty() {
        rootNode.showPretty(indent: 0)
    }
}

extension Trie where Element == String, Path: DecomposedStringProtocol {
    convenience init(_ strings: [String]) {
        self.init()
        insertBatch(strings)
    }

    func insertBatch(_ strings: [String]) {
        strings.forEach { string in
            insert(string)
        }
    }

    func insert(_ element: Element) {
        self.insert(element, at: Path.decompose(element))
    }

    func find(from string: String) -> [Element] {
        return self.find(from: Path.decompose(string))
    }

    func find(prefix: String) -> [Element] {
        return self.find(from: Path.decompose(prefix))
    }
}

protocol DecomposedStringProtocol: Hashable {
    associatedtype C: Collection where C.Element == Self
    static func decompose(_ string: String) -> C
}

extension UTF16.CodeUnit: DecomposedStringProtocol {
    static func decompose(_ string: String) -> String.UTF16View {
        string.utf16
    }
}

extension UTF8.CodeUnit: DecomposedStringProtocol {
    static func decompose(_ string: String) -> String.UTF8View {
        string.utf8
    }
}

extension Character: DecomposedStringProtocol {
    static func decompose(_ string: String) -> String {
        string
    }
}

let trie = Trie<String, UTF16.CodeUnit>()
trie.insert("CEZANNE")
trie.insert("CANMAKE")
trie.insert("Dior")
trie.insert("ELIXIR")
trie.showPretty()


print(trie.find(from: "fffCE"))
print(trie.find(prefix: "C"))