Cod sursa(job #3369252)

Utilizator DavidLDavid Lauran DavidL Data 18 octombrie 2026 16:53:41
Problema Aho-Corasick Scor Ascuns
Compilator cpp-64 Status waiting
Runda Arhiva educationala Marime 9.05 kb
#include <bits/stdc++.h>
using namespace std;

ifstream fin("ahocorasick.in");
ofstream fout("ahocorasick.out");

class TrieNode {
public:
    static const int ALPHABET_SIZE = 26;

private:
    char letter;
    vector<int> wordsEndingHereIndexes;
    int countMatched;
    TrieNode* root;
    TrieNode* parent;
    TrieNode *children[ALPHABET_SIZE];
    TrieNode *longestSubfix;

public:
    TrieNode(char letter) : letter{ letter }, countMatched{ 0 }, parent{ nullptr } , longestSubfix { nullptr } {
        for (int i = 0; i < ALPHABET_SIZE; i++) {
            children[i] = nullptr;
        }
        root = this;
    }
    TrieNode(TrieNode* root, char letter) : letter{ letter }, countMatched{ 0 }, parent{ nullptr }, longestSubfix { nullptr } {
        for (int i = 0; i < ALPHABET_SIZE; i++) {
            children[i] = nullptr;
        }
        this->root = root;
    }
    ~TrieNode() {
        for (int i = 0; i < ALPHABET_SIZE; i++) {
            if (children[i] != nullptr) {
                delete children[i];
            }
        }
    }

    char getLetter() const { return letter; }

    TrieNode* getParent() const { return parent; }
    void setParent(TrieNode* parent) {
        this->parent = parent;
    }

    bool hasChild(char childLetter) const {
        // cout << "haschildstart. childlettecoder=" << letterToIntCode(childLetter) << "\n";
        // cout << "is child null? " << (children[letterToIntCode(childLetter)] == nullptr) << "\n";
        bool toReturn = children[letterToIntCode(childLetter)] != nullptr; 
        // cout << "haschildend\n";
        return toReturn;
    }
    TrieNode* getChild(char childLetter) const {
        return children[letterToIntCode(childLetter)];
    }

    // Precondition: the child doesn't exist yet
    void addChild(char childLetter) {
        int childLetterIntCode = letterToIntCode(childLetter); 
        children[childLetterIntCode] = new TrieNode(root, childLetter);
        children[childLetterIntCode]->setParent(this);
    }

    TrieNode* getLongestSubfix() { return longestSubfix; }
    void setLongestSubfix(TrieNode* longestSubfix) {
        this->longestSubfix = longestSubfix;
    }

    void addWordEndingHere(int wordIndex) {
        wordsEndingHereIndexes.push_back(wordIndex);
    }
    int getCountWordsEndingHere() const { return wordsEndingHereIndexes.size(); }
    const vector<int>& getWordsEndingHere() const { return wordsEndingHereIndexes; }

    void addMatch() {
        countMatched++;
    }
    void addMatches(int newMatchesCount) {
        countMatched += newMatchesCount;
    }
    int getCountMatched() const { return countMatched; }

    bool isRoot() const { 
        // cout << "enter isroot\n";
        bool toReturn = (root == this);
        // cout << "exit isroot with " << toReturn << "\n";
        return toReturn;
    }
    TrieNode* getRoot() const { return root; }

private:
    static inline int const letterToIntCode(char letter) {
        return letter - 'a';
    }
};

class Trie {
public:
    static const char ROOT_LETTER = '0'; 

private:
    TrieNode* root;

public:
    Trie() {
        root = new TrieNode(ROOT_LETTER);
    }
    ~Trie() { delete root; }

    void addWord(string word, int wordIndex) {
        addWord(root, word, 0, wordIndex);
    }
    
    TrieNode* getRoot() const {
        return root;
    }

    vector<TrieNode*> getNodesOrderedByDepth() const {
        vector<TrieNode*> nodesOrderedByDepth;
        nodesOrderedByDepth.push_back(root);
        int currentNodeIndex = 0;

        while (currentNodeIndex < (int)nodesOrderedByDepth.size()) {
            const TrieNode* currentNode = nodesOrderedByDepth[currentNodeIndex];
            // cout << "currentnode==null? " << (currentNode == nullptr) << "\n";

            for (char childLetter = 'a'; childLetter <= 'z'; childLetter++) {
                if (currentNode->hasChild(childLetter)) {
                    nodesOrderedByDepth.push_back(currentNode->getChild(childLetter));
                }
            }

            currentNodeIndex++;
        }

        return nodesOrderedByDepth;
    }

private:
    void addWord(TrieNode* currentNode, const string& word, int nextLetterIndex, int wordIndex) {
        // cout << "addword currentnode==nullptr? " << (currentNode == nullptr) << "\n";
        if (nextLetterIndex == (int)word.size()) {
            // cout << "adding word having size " << word.size() << ". is currentnode root? " << (currentNode->isRoot()) << "\n";
            currentNode->addWordEndingHere(wordIndex);
            return;
        }

        char nextLetter = word[nextLetterIndex];
        if (!currentNode->hasChild(nextLetter)) {
            currentNode->addChild(nextLetter);
        }
        addWord(currentNode->getChild(nextLetter), word, nextLetterIndex + 1, wordIndex);
    }
};

void computeLongestSubfixes(Trie* trie) {
    // BFS + recursive formula

    vector<TrieNode*> nodesOrderedByDepth = trie->getNodesOrderedByDepth();

    for (TrieNode* currentNode : nodesOrderedByDepth) {
        // cout << "iscurrentnode nullptr? " << (currentNode == nullptr) << "\n";
        if (currentNode->isRoot()) {
            continue;
        }

        TrieNode* longestSubfixParent = currentNode->getParent()->getLongestSubfix();

        while (longestSubfixParent != nullptr && !longestSubfixParent->hasChild(currentNode->getLetter())) {
            longestSubfixParent = longestSubfixParent->getLongestSubfix();
        }
        if (longestSubfixParent == nullptr) {
            longestSubfixParent = trie->getRoot();
        }
        
        if (longestSubfixParent->hasChild(currentNode->getLetter()) && longestSubfixParent->getChild(currentNode->getLetter()) != currentNode) {
            currentNode->setLongestSubfix(longestSubfixParent->getChild(currentNode->getLetter()));
        } else {
            currentNode->setLongestSubfix(nullptr);
        }
        // cout << "longestsubfix of currentnode " << currentNode->getLetter() << " which does have a word? " <<
        // (currentNode->getCountWordsEndingHere() > 0) << "is the node " << (currentNode->getLongestSubfix() == nullptr ? 'xxx' : currentNode->getLongestSubfix()->getLetter()) << " which does have a word? " << (currentNode->getLongestSubfix() != nullptr && currentNode->getLongestSubfix()->getCountWordsEndingHere() > 0) << "\n";
    }
}

void addCountMatchedToLongestSubfixBottomUp(Trie* dictionaryTrie, const vector<TrieNode*>& nodesOrderedByDepth) {
    for (int nodeIndex = nodesOrderedByDepth.size() - 1; nodeIndex >= 0; nodeIndex--) {
        TrieNode* currentNode = nodesOrderedByDepth[nodeIndex];
        TrieNode* longestSubfix = currentNode->getLongestSubfix();

        if (longestSubfix != nullptr) {
            longestSubfix->addMatches(currentNode->getCountMatched());
        }
    }
}

vector<int> searchWords(Trie* dictionaryTrie, const string& searchString, int countWords) {
    // searchWords(dictionaryTrie->getRoot(), dictionaryTrie->getRoot(), searchString);
    TrieNode* currentNode = dictionaryTrie->getRoot();
    for (int nextLetterIndex = 0; nextLetterIndex < (int)searchString.size(); nextLetterIndex++) {
        char nextLetter = searchString[nextLetterIndex];
        while (currentNode != nullptr && !currentNode->isRoot() && !currentNode->hasChild(nextLetter)) {
            currentNode = currentNode->getLongestSubfix();
        }
        if (currentNode == nullptr) {
            currentNode = dictionaryTrie->getRoot();
        }

        currentNode = currentNode->hasChild(nextLetter) ? currentNode->getChild(nextLetter) : dictionaryTrie->getRoot();
        currentNode->addMatch();
    }
    
    // step 4: sum up the findings (if dotted line from node x to node y, then add the value of node x to node y)
    vector<TrieNode*> nodesOrderedByDepth = dictionaryTrie->getNodesOrderedByDepth();
    addCountMatchedToLongestSubfixBottomUp(dictionaryTrie, nodesOrderedByDepth);
    
    vector<int> occurrences(countWords);
    for (TrieNode* node : nodesOrderedByDepth) {
        for (int wordIndex : node->getWordsEndingHere()) {
            occurrences[wordIndex] = node->getCountMatched();
        }
    }

    return occurrences;
}

void print(const vector<int>& toPrint) {
    for (int element : toPrint) {
        fout << element << "\n";
    }
}

int main() {
    // step 1: build the Trie with its normal edges
    // step 2: figure out the dotted (longest subfix) edges
    // step 3: simple search thru the Trie
    // step 4: sum up the findings (if dotted line from node x to node y, then add the value of node x to node y)

    string largeString;
    getline(fin, largeString);
    // cout << "lagestring=" << largeString << "\n";

    Trie* dictionaryTrie = new Trie();

    int dictionaryLength;
    fin >> dictionaryLength;

    string unused;
    getline(fin, unused);

    for (int i = 0; i < dictionaryLength; i++) {
        string word;
        getline(fin, word);
        // cout << "read word " << word << "\n";
        dictionaryTrie->addWord(word, i);
    }

    computeLongestSubfixes(dictionaryTrie);
    // exit(0);

    vector<int> countOccurrences = searchWords(dictionaryTrie, largeString, dictionaryLength);
    print(countOccurrences);

    delete dictionaryTrie;

    return 0;
}