7
\$\begingroup\$

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();
}
\$\endgroup\$
  • \$\begingroup\$ Can you give more explanation about all of these xxxxDFS functions please? \$\endgroup\$ – Calak Nov 19 '18 at 22:19
  • \$\begingroup\$ All the DFS (depth first search) function are basically the recursive functions to go down the tree and do something to each node. For example, deleteTrieDFS traverses each node's children recursively until the node doesn't have children - then the node is deleted (set to nullptr). Does it make sense? \$\endgroup\$ – user_185051 Nov 19 '18 at 22:26
5
+50
\$\begingroup\$

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>;
};
| improve this answer | |
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.