30 Day Challenge Day 16 | Leetcode 701. Insert into a Binary Search Tree

题解

Hard

方法：Trie

struct TrieNode {
    vector<int> indexes;
    vector<TrieNode*> next;
    
    TrieNode() : next(26, nullptr) {
        
    }
};

class Solution {
public:
    vector<vector<string>> wordSquares(vector<string>& words) {
        // the size of a word determine the square size
        int num = words.size(), size = words[0].size();

        vector<vector<string>> results;
        vector<string> result;
        
        TrieNode* root = buildTrieTree(words);
        
        for(auto word : words) {
            result.push_back(word);
            search(words, size, root, result, results);
            result.pop_back();
        }
        
        return results;
    }
    
    TrieNode* buildTrieTree(vector<string>& words) {
        TrieNode* root = new TrieNode();

        for(int i = 0; i < words.size(); i++) {
            TrieNode* node = root;
            for(auto c : words[i]) {
                if(!(node->next[c-'a'])) node->next[c-'a'] = new TrieNode();
                node = node->next[c-'a'];
                node->indexes.push_back(i);
            }
        }
        
        return root;
    }
    
    void search(vector<string>& words, int size, TrieNode* root, vector<string> result, vector<vector<string>>& results) {
        if(result.size() == size) {
            results.push_back(result);
            return;
        }
        
        string prefix;
        for(int i = 0; i < result.size(); i++) {
            prefix += result[i][result.size()];
        }
        
        TrieNode* node = root;
        for(auto c : prefix) {
            if(!(node->next[c-'a'])) return;
            node = node->next[c-'a'];
        }
        
        for(int i : node->indexes) {
            result.push_back(words[i]);
            search(words, size, root, result, results);
            result.pop_back();
        }
    }
};