Trie implementation for strings in Swift

Question

A trie for handling strings for an autocomplete dictionary. This passes my fairly casual tests, though it's always possible that there are broken edge cases, but I'm mainly concerned about design and efficiency: is this a sensible way to implement this data structure? (In particular, is it sensible to do the insert method recursively?). It feels a little gnarly to me.

struct Trie {
    var root: Node
    init() {
        root = Node()
    }
    func search(_ word: String) -> Bool {
        let letters = Array(word)
        var curnode = root
        for letter in letters {
            guard let match = curnode.children.first(where: {(key, _) in
                key == letter
            })
                else {
                    return false
            }
            curnode = match.value
        }
        if curnode.contained {
            return true
        }
        return false
    }
    
    func remove(_ word: String) {
        let letters = Array(word)
        var curnode = root
        for letter in letters {
            if !curnode.children.contains(where: {(key, _) in
                key == letter
            }) {
                break
            } else {
                curnode = curnode.children[letter]!
            }
        }
        curnode.contained = false
    }
    
    func insert(_ letters: [Character], parent: Node) -> Node {
        if letters.count == 1 {
            let letter = letters[0]
            if parent.children.contains(where: {(key, _) in
                key == letter
            }) {
                let newNode = parent
                newNode.children[letter]!.contained = true
                return newNode
            } else {
                let newNode = Node(letter, final: true)
                return newNode
            }
        } else {
            let first = letters[0]
            let rest = Array(letters.dropFirst())
            if let subtree = parent.children.first(where: {(key, _) in
                           key == first
            }) {
                let newNode = Node(first, final: subtree.value.contained, kids: subtree.value.children)
                newNode.children[rest[0]] = insert(rest, parent: newNode)
                return newNode

            } else {
            let newNode = Node(first, final: false)
            newNode.children[rest[0]] = insert(rest, parent: newNode)
            return newNode
            }
        }
    }
    mutating func insert(_ word: String) {
        let new_subtree = insert(Array(word), parent: root)
        root.children[new_subtree.char!] = new_subtree
    }
}

class Node {
    var char: Character?
    var children: [Character:Node]
    var contained: Bool
    init() {
        char = nil
        children = [:]
        contained = false
    }
    init(_ c: Character, final: Bool) {
        children = [:]
        contained = final
        char = c
    }
    init(_ c: Character, final: Bool, kids: [Character:Node]) {
        children = kids
        contained = final
        char = c
    }
}

Clarification edit: this isn't meant to be production ready, is only meant to be a concise and straightforward implementation of the data structure. So I intentionally left off checks for problem input; it should only be able to handle sensible inputs right now, where "sensible" means "input that someone using the data structure correctly to add, look for, and remove data would intentionally provide." So, not empty strings, removing stuff that isn't there, inserting the same thing twice, etc. etc.

The casual test cases I used, meant to represent sensible behavior, are:

var testTrie = Trie()

testTrie.insert("cat")
testTrie.search("cat") // T
testTrie.search("car") // F
testTrie.insert("car")
testTrie.search("car") // T
testTrie.search("cat") // T
testTrie.search("ca") // F
testTrie.search("cad") // F
testTrie.search("carburetor") // F
testTrie.insert("carburetor")
testTrie.search("carburetor") // T
testTrie.search("car") // T
testTrie.search("cat") // T
testTrie.search("ca") // F    
testTrie.remove("car")
testTrie.search("carburetor") // T
testTrie.search("car") // F
testTrie.search("cat") // T

Answer 1

String traversal

At several places in your code you convert a string to an array of its characters in order to iterate over it:

let letters = Array(word)
for letter in letters {
    // ...
}

These intermediate arrays are not needed. A Swift string is a collection of characters, so that you can iterate over it simply with

for letter in word {
    // ...
}

Dictionary access

In the search() method you look up the node for a character with

guard let match = curnode.children.first(where: {(key, _) in
    key == letter
})
    else {
        return false
}
curnode = match.value

Similar patterns are also in the other methods. This dictionary lookup can be simplified using a subscript:

guard let node = curnode.children[letter] else {
    return false
}
curnode = node

Returning boolean values

Code like

if someCondition {
     return true
}
return false

can always be simplified to

return someCondition

which is shorter and clearer. The search method then looks like this:

func search(_ word: String) -> Bool {
    var curnode = root
    for letter in word {
        guard let node = curnode.children[letter] else {
            return false
        }
        curnode = node
    }
    return curnode.contained
}

Removing non-existent strings

Removing a string which has never been inserted currently has undesired side effects:

var trie = Trie()
trie.insert("a")
trie.remove("ab")
print(trie.search("a")) // false

That is easy to fix: As soon as the traversal does not find a node for the next character it should return instead of setting curnode.contained = false on the last node encountered:

func remove(_ word: String) {
    var curnode = root
    for letter in word {
        guard let node = curnode.children[letter] else {
            return // <--- HERE
        }
        curnode = node
    }
   curnode.contained = false
}

Mutating (or not?) methods

The mutating keyword in

mutating func insert(_ word: String)

is not needed: Node is a reference type so that the properties of root can be modified without making the method mutating. For the same reason, the property can be declared as a constant:

struct Trie {
    let root: Node
    // ...
}

Use substrings!

The main insert method create a array of all characters:

 let new_subtree = insert(Array(word), parent: root)

and the recursive helper methods repeatedly creates more arrays of the remaining characters:

 let rest = Array(letters.dropFirst())

That is very inefficient. The better approach is that the helper method takes a Substring argument:

func insert(_ letters: Substring, parent: Node) -> Node

so that it can call itself with

let rest = letters.dropFirst()
insert(rest, parent: newNode)

This is called “slicing” in Swift and very efficient because the substrings share the element storage with the original string and no copies are made.

The main insert method then calls the helper method with a substring comprising all its characters:

func insert(_ word: String) {
    let new_subtree = insert(word[...], parent: root)
    // ...
}

Simplify the insertion method (and the Node type)

I found the insertion code difficult to understand. It also has some problems (which you are already aware of):

• It is not possible to insert an empty string.
• It is not possible to insert the same string twice.

To be honest: I cannot see which cases are handled correctly and which are not.

What I also do not like is the var char: Character? property of Node. Apparently this is needed to insert a newly created subtree at the right position of the parent's children dictionary. But

• it introduces some redundancy,
• it is not clear in which cases it can be nil (only in the root node?),
• it requires forced unwrapping.

Doing the insertion recursively is fine. But if we create new nodes before the recursive call with the rest of the string then everything becomes much simpler:

func insert(_ word: Substring, node: Node) {
    if let letter = word.first {
        if let nextnode = node.children[letter] {
            insert(word.dropFirst(), node: nextnode)
        } else {
            let newnode = Node()
            node.children[letter] = newnode
            insert(word.dropFirst(), node: newnode)
        }
    } else {
        node.contained = true
    }
}

func insert(_ word: String) {
    insert(word[...], node: root)
}

The char property is not needed anymore, i.e. that type simplifies to

class Node {
    var children: [Character: Node] = [:]
    var contained: Bool = false
}

More advantages:

• The recursion terminates when the string is empty, not when it is a single character. As a consequence, inserting an empty string works now.
• Inserting the same string twice works as well.

The same can be done with iteration instead of recursion:

func insert(_ word: String) {
    var curnode = root
    for letter in word {
        if let nextnode = curnode.children[letter] {
            curnode = nextnode
        } else {
            let newnode = Node()
            curnode.children[letter] = newnode
            curnode = newnode
        }
    }
    curnode.contained = true
}

That is a matter of taste, but it is shorter and makes even the substrings obsolete.

Naming

You use different naming conventions in your code:

curnode, newNode, new_subtree

The Swift naming convention is camelcase (upper camelcase for types, and lower camelcase for everything else):

currentNode, newNode, newSubtree

I would also prefer char or character over letter: A Swift string can contain arbitrary Unicode characters, not only “letters.”