Trie (dictionary tree) data structure implementation: insertWord, printAllWords, searchPrefix, deleteWord, deleteTrie

Question

This Trie implementation works, however, I would like to get an advice on how to improve this code.

Any advice is appreciated: functions implementation, memory management, modern C++ usage, const-correctness, code readability, anything else I am missing? Any additional functions that would be useful?

Trie.hpp

#pragma once

#include <iostream>
#include <string>
#include <vector>

/*
 Trie for words, which consist of
 lowercase English letters only
 */

class Trie {

private:
    struct TrieNode {
        bool isEndOfWord;
        std::vector<std::shared_ptr<TrieNode>> children;
        TrieNode() : children(26), isEndOfWord(false) {}
    };

    std::shared_ptr<TrieNode> root = std::make_shared<TrieNode>();

    bool hasChildren (const std::shared_ptr<TrieNode>& root) const;
    void printDFS (const std::shared_ptr<TrieNode>& root, std::string& stringToPrint) const;
    void searchPrefix (const std::shared_ptr<TrieNode>& root, const std::string& prefix) const;
    void searchPrefixDFS (const std::shared_ptr<TrieNode>& root, std::string& currPrefix) const;
    bool deleteWordDFS (std::shared_ptr<TrieNode>& root, const std::string& word) const;
    void deleteTrieDFS (std::shared_ptr<TrieNode>& root);

public:
    Trie();
    void insertWord (const std::string& word);
    void printAllWords () const;
    void searchPrefix (const std::string& prefix);
    bool deleteWord (const std::string& word);
    void deleteTrie ();
    ~Trie();
};

Trie.cpp

#include "Trie.hpp"

Trie::Trie () {
    // Should I be initializing anything here?
}

void Trie::insertWord (const std::string& word) {
    std::shared_ptr<TrieNode> curr = root;
    for (const auto& ch : word) {
        if (!curr->children[ch - 'a']) {
            //std::cout << "adding: " << (char)(ind + 'a') << "\n";
            curr->children[ch - 'a'] = std::make_shared<TrieNode>();
        }
        curr = curr->children[ch - 'a'];
    }
    curr->isEndOfWord = true;
}

bool Trie::hasChildren (const std::shared_ptr<TrieNode>& root) const {
    for (const auto& child : root->children) {
        if (child) {
            return true;
        }
    }
    return false;
}

void Trie::printDFS (const std::shared_ptr<TrieNode>& root,
                     std::string& stringToPrint) const {
    if (root->isEndOfWord) {
        std::cout << stringToPrint << std::endl;
    }
    for (int i = 0; i < 26; ++i) {
        if (root->children[i]) {
            stringToPrint.push_back(i + 'a');
            printDFS(root->children[i], stringToPrint);
            stringToPrint.pop_back();
        }
    }
}

void Trie::printAllWords () const {
    std::string stringToPrint;
    if (root) {
        Trie::printDFS (root, stringToPrint);
    }
}

void Trie::searchPrefixDFS (const std::shared_ptr<TrieNode>& root, std::string& currPrefix) const {
    if (root->isEndOfWord) {
        std::cout << currPrefix << "\n";
    }
    if (!hasChildren(root)){
        return;
    }
    for (int i = 0; i < 26; ++i) {
        if (root->children[i]) {
            currPrefix.push_back(i + 'a');
            searchPrefixDFS(root->children[i], currPrefix);
        }
    }
}

void Trie::searchPrefix (const std::shared_ptr<TrieNode>& root, const std::string& prefix) const {
    std::shared_ptr<TrieNode> curr = root;
    for (int lev = 0; lev < prefix.length(); lev++) {
        char currChar = prefix[lev];
        if (!curr->children[currChar - 'a']) {
            std::cout << "Prefix " << prefix << " is not in a trie.\n";
            return;
        }
        curr = curr->children[currChar - 'a'];
    }
    std::string currPrefix = prefix;
    Trie::searchPrefixDFS(curr, currPrefix);
}

void Trie::searchPrefix (const std::string& prefix) {
    if (prefix.empty()) {
        Trie::printAllWords();
    } else {
        Trie::searchPrefix (root, prefix);
    }
}

bool Trie::deleteWordDFS (std::shared_ptr<TrieNode>& curr,
                          const std::string& word) const {

    // This only returns true if we need ot keep going
    // This returns false if we are finished with deleting the word
    // or if the word in not in the trie

    if (!curr) {
        return false;
    }
    if (word.length()) {
        if (curr && curr->children[word[0] - 'a'] &&
            deleteWordDFS(curr->children[word[0] - 'a'], word.substr(1)) &&
            !curr->isEndOfWord) {
            if (!hasChildren(curr)) {
                curr = nullptr;
                return true;
            } else {
                return false;
            }
        }
    }
    if (word.length() == 0 && curr->isEndOfWord) {
        // end of word
        if (!hasChildren(curr)) {
            // if no children - delete it
            curr = nullptr;
            return true; // keep going to delete parents
        } else {
            curr->isEndOfWord = false;
            return false; // do not keep going
        }
    }
    return false;
}

bool Trie::deleteWord (std::string const & word) {
    return Trie::deleteWordDFS(root, word);
}

void Trie::deleteTrieDFS (std::shared_ptr<TrieNode>& root) {
    for (auto& node : root->children) {
        if (node) {
            deleteTrieDFS(node);
        }
    }
    if (!hasChildren(root)) {
        root = nullptr;
    }
}

void Trie::deleteTrie () {
    Trie::deleteTrieDFS(root);
}

Trie::~Trie () {
    // Should I be deleting anything here?
}

main.cpp

#include "Trie.hpp"

int main() {

    // Vector of words
    std::vector<std::string> words;
    words.push_back("why");
    words.push_back("a");
    words.push_back("peach");
    words.push_back("app");
    words.push_back("cat");
    words.push_back("apple");
    words.push_back("pea");
    words.push_back("banana");
    words.push_back("ban");

    // Trie init
    Trie trieDict;

    // Add words to trie
    for (const auto& word : words) {
        trieDict.insertWord(word);
    }

    std::cout << "Print all trie words: \n";
    trieDict.printAllWords();

    // Search prefix:
    std::string prefix = "p";
    std::cout << "Words starting with " << prefix << ":\n";
    trieDict.searchPrefix(prefix);

    // Deleting words from the trie
    trieDict.deleteWord("why");
    trieDict.deleteWord("app");
    //trieDict.deleteWord("a");
    //trieDict.deleteWord("peach");
    //trieDict.deleteWord("pea");
    //trieDict.deleteWord("cat");
    //trieDict.deleteWord("ban");

    std::cout << "Print all trie words after deletion: \n";
    trieDict.printAllWords();

    // Delete Trie
    trieDict.deleteTrie();
    std::cout << "Print trie after trie deletion: \n";
    trieDict.printAllWords();
}

Answer 1

Not an exhaustive review but a few remarks about

struct TrieNode

You are using an std::vector<T>, although its size is fixed at 26. Consider using an std::array<T, 26> instead. This reduces calls to allocate dynamic memory. Also, your loops like

for (int i = 0; i < 26; ++i) {
    if (root->children[i]) {
        //do something
    }
}

are then guaranteed to not cause undefined behaviour due to illegal access. Additionally, since every node has at most one parent, there is no reason to use std::shared_ptr, go for std::unique_ptr instead.

26 as the size of your alphabet is a magic number (explanation on wikipedia, stackoverflow), although a rather obvious one. Nonetheless, you could name it, e.g. declare a

static constexpr size_t alphabet_size = 26;

One could consider a handy using for later on (std::shared_ptr<TrieNode> is repeated quite often in your code) as well as default values, which altogether results in something along the lines of

class Trie {
private:
    static constexpr size_t alphabet_size = 26;
    struct TrieNode {
        bool isEndOfWord = false;
        std::array<std::unique_ptr<TrieNode>, alphabet_size> children = {};
    };

    using TrieNodePtr = std::unique_ptr<TrieNode>;
};