Advertisement
maxim_shlyahtin

Aho-Corasik mod

May 30th, 2024
694
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 5.32 KB | None | 0 0
  1. #include <iostream>
  2. #include <map>
  3. #include <vector>
  4. #include <string>
  5. #include <array>
  6. #include <utility>
  7. #include <algorithm>
  8.  
  9. const int N = 5;
  10. const char WILDCARD = '?';
  11.  
  12. std::vector<std::string> patterns;
  13. std::array<char, N> alphabet = { 'A', 'C', 'T', 'G', 'N' };
  14.  
  15. struct Node {
  16.     std::map<char, Node*> son;
  17.     std::map<char, Node*> go;
  18.     Node* parent;
  19.     Node* suffLink;
  20.     Node* up;
  21.     char charToParent;
  22.     bool isTerminal;
  23.     std::vector<int> leafPatternNumber;
  24.  
  25.     Node() {
  26.         son = { {alphabet[0], NULL}, {alphabet[1], NULL}, {alphabet[2], NULL}, {alphabet[3], NULL}, {alphabet[4], NULL}, };
  27.         go = { {alphabet[0], NULL}, {alphabet[1], NULL}, {alphabet[2], NULL}, {alphabet[3], NULL}, {alphabet[4], NULL}, };
  28.         parent = nullptr;
  29.         suffLink = nullptr;
  30.         up = nullptr;
  31.         charToParent = ' ';
  32.         isTerminal = false;
  33.     }
  34. };
  35.  
  36. struct Trie {
  37.     Node* root;
  38.  
  39.     Trie() {
  40.         root = new Node();
  41.     }
  42.  
  43.     Node* getSuffLink(Node* v) {
  44.         if (v->suffLink == nullptr) {
  45.             if (v == root || v->parent == root) {
  46.                 v->suffLink = root;
  47.             } else {
  48.                 v->suffLink = getLink(getSuffLink(v->parent), v->charToParent);
  49.             }
  50.         }
  51.         return v->suffLink;
  52.     }
  53.  
  54.     Node* getLink(Node* v, char c) {
  55.         if (v->go[c] == nullptr) {
  56.             if (v->son[c]) {
  57.                 v->go[c] = v->son[c];
  58.             } else if (v == root) {
  59.                 v->go[c] = root;
  60.             } else {
  61.                 v->go[c] = getLink(getSuffLink(v), c);
  62.             }
  63.         }
  64.         return v->go[c];
  65.     }
  66.  
  67.     Node* getUp(Node* v) {
  68.         if (v->up == nullptr) {
  69.             if (getSuffLink(v)->isTerminal) {
  70.                 v->up = getSuffLink(v);
  71.             } else if (getSuffLink(v) == root) {
  72.                 v->up = root;
  73.             } else {
  74.                 v->up = getUp(getSuffLink(v));
  75.             }
  76.         }
  77.         return v->up;
  78.     }
  79.  
  80.     void addString(const std::string& s, int patternNumber) {
  81.         Node* curr = root;
  82.         for (char c : s) {
  83.             if (curr->son[c] == nullptr) {
  84.                 curr->son[c] = new Node();
  85.                 curr->son[c]->parent = curr;
  86.                 curr->son[c]->charToParent = c;
  87.             }
  88.             curr = curr->son[c];
  89.         }
  90.         curr->isTerminal = true;
  91.         curr->leafPatternNumber.push_back(patternNumber);
  92.     }
  93.  
  94.     void processText(const std::string& text, std::vector<std::pair<int, int>>& results) {
  95.         Node* curr = root;
  96.         for (size_t i = 0; i < text.length(); ++i) {
  97.             char c = text[i];
  98.             curr = getLink(curr, c);
  99.             Node* checkNode = curr;
  100.             while (checkNode != root) {
  101.                 if (checkNode->isTerminal) {
  102.                     for (int patNum : checkNode->leafPatternNumber) {
  103.                         results.push_back({ i + 2 - patterns[patNum].length(), patNum + 1 });
  104.                     }
  105.                 }
  106.                 checkNode = getUp(checkNode);
  107.             }
  108.         }
  109.     }
  110. };
  111.  
  112. std::vector<int> findPatternWithWildcards(const std::string& text, const std::string& pattern, char wildcard) {
  113.     std::vector<std::string> subPatterns;
  114.     std::vector<int> startPositions;
  115.     std::string currentSubPattern;
  116.     for (size_t i = 0; i < pattern.length(); ++i) {
  117.         if (pattern[i] == wildcard) {
  118.             if (!currentSubPattern.empty()) {
  119.                 subPatterns.push_back(currentSubPattern);
  120.                 startPositions.push_back(i - currentSubPattern.length() + 1);
  121.                 currentSubPattern.clear();
  122.             }
  123.         } else {
  124.             currentSubPattern += pattern[i];
  125.         }
  126.     }
  127.     if (!currentSubPattern.empty()) {
  128.         subPatterns.push_back(currentSubPattern);
  129.         startPositions.push_back(pattern.length() - currentSubPattern.length() + 1);
  130.     }
  131.  
  132.     Trie trie;
  133.     for (size_t i = 0; i < subPatterns.size(); ++i) {
  134.         trie.addString(subPatterns[i], i);
  135.     }
  136.  
  137.     for (const auto& it : subPatterns) {
  138.         patterns.push_back(it);
  139.     }
  140.  
  141.     std::vector<std::pair<int, int>> results;
  142.     trie.processText(text, results);
  143.  
  144.     std::vector<int> C(text.length(), 0);
  145.     for (const auto& result : results) {
  146.         int pos = result.first - startPositions[result.second - 1];
  147.         if (pos >= 0 && pos + pattern.length() <= text.length()) {
  148.             C[pos]++;
  149.         }
  150.     }
  151.  
  152.     std::vector<int> finalPositions;
  153.     for (size_t i = 0; i < C.size(); ++i) {
  154.         if (C[i] == subPatterns.size()) {
  155.             finalPositions.push_back(i);
  156.         }
  157.     }
  158.  
  159.     // Фильтрация пересекающихся вхождений
  160.     std::vector<int> nonOverlappingPositions;
  161.     int lastEnd = -1;
  162.     for (int pos : finalPositions) {
  163.         if (pos > lastEnd) {
  164.             nonOverlappingPositions.push_back(pos);
  165.             lastEnd = pos + pattern.length() - 1;
  166.         }
  167.     }
  168.  
  169.     return nonOverlappingPositions;
  170. }
  171.  
  172. int main() {
  173.     std::string text;
  174.     std::string pattern;
  175.     char wildcard;
  176.     std::cin >> text >> pattern >> wildcard;
  177.  
  178.     std::vector<int> positions = findPatternWithWildcards(text, pattern, wildcard);
  179.     for (int pos : positions) {
  180.         std::cout << pos << std::endl;
  181.     }
  182.  
  183.     return 0;
  184. }
  185.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement