fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int INF = 1e9;
  5.  
  6. // ===================================================================
  7. // 1) Core SAM structure and construction
  8. //    This is the foundation. All other functions use this class.
  9. // ===================================================================
  10.  
  11. // 1.1) SuffixAutomaton class
  12. //      Purpose:
  13. //        - Builds a SAM for a given string in O(n) time.
  14. //        - Stores all necessary arrays (len, link, next) for the automaton.
  15. //      Input:
  16. //        - s : the input string (typically lowercase letters).
  17. //      Output:
  18. //        - An object that contains the SAM.
  19. //      Time complexity:
  20. //        - O(n * alphabet_size) with array transitions.
  21. //      Constraints:
  22. //        - The string length n must be known.
  23. //        - The alphabet size is assumed constant (e.g., 26 for lowercase).
  24. //      Notes:
  25. //        - The number of states is at most 2*n.
  26. //        - The root is state 0.
  27. //        - The `extend` function adds one character at a time.
  28. //        - This class is used as a black box by all other functions.
  29. struct SuffixAutomaton {
  30.     static const int ALPHABET = 26;          // for lowercase English letters
  31.     vector<array<int, ALPHABET>> next;       // transitions
  32.     vector<int> link;                        // suffix links
  33.     vector<int> len;                         // longest length of strings in this state
  34.     int last;                                // state corresponding to the whole current string
  35.  
  36.     // Extra members for advanced queries
  37.     vector<bool> isClone;                    // true if state is a clone
  38.     vector<int> occ;                         // occurrence count (will be computed)
  39.     vector<int> firstPos;                    // earliest end position of substrings in this state
  40.     vector<int> lastPos;                     // latest end position of substrings in this state
  41.  
  42.     SuffixAutomaton() {}
  43.     SuffixAutomaton(const string& s) {
  44.         init(s);
  45.     }
  46.  
  47.     void init(const string& s) {
  48.         next.clear();
  49.         link.clear();
  50.         len.clear();
  51.         isClone.clear();
  52.         occ.clear();
  53.         firstPos.clear();
  54.         lastPos.clear();
  55.  
  56.         // Root (state 0)
  57.         next.push_back({});
  58.         for (auto &a : next[0]) a = -1;
  59.         link.push_back(-1);
  60.         len.push_back(0);
  61.         isClone.push_back(false);
  62.         occ.push_back(0);
  63.         firstPos.push_back(INF);
  64.         lastPos.push_back(-INF);
  65.         last = 0;
  66.  
  67.         for (char c : s) {
  68.             extend(c - 'a');
  69.         }
  70.     }
  71.  
  72.     void extend(int c) {
  73.         int cur = (int)next.size();
  74.         next.push_back({});
  75.         for (auto &a : next.back()) a = -1;
  76.         len.push_back(len[last] + 1);
  77.         link.push_back(0);
  78.         isClone.push_back(false);
  79.         occ.push_back(1);                     // new state represents a prefix
  80.         firstPos.push_back(len.back() - 1);   // end position of this prefix
  81.         lastPos.push_back(len.back() - 1);
  82.  
  83.         int p = last;
  84.         while (p != -1 && next[p][c] == -1) {
  85.             next[p][c] = cur;
  86.             p = link[p];
  87.         }
  88.         if (p == -1) {
  89.             link[cur] = 0;
  90.         } else {
  91.             int q = next[p][c];
  92.             if (len[p] + 1 == len[q]) {
  93.                 link[cur] = q;
  94.             } else {
  95.                 int clone = (int)next.size();
  96.                 next.push_back(next[q]);      // copy transitions
  97.                 len.push_back(len[p] + 1);
  98.                 link.push_back(link[q]);
  99.                 isClone.push_back(true);
  100.                 occ.push_back(0);
  101.                 firstPos.push_back(INF);
  102.                 lastPos.push_back(-INF);
  103.  
  104.                 while (p != -1 && next[p][c] == q) {
  105.                     next[p][c] = clone;
  106.                     p = link[p];
  107.                 }
  108.                 link[q] = link[cur] = clone;
  109.             }
  110.         }
  111.         last = cur;
  112.     }
  113.  
  114.     // Returns the number of states (including root).
  115.     int size() const {
  116.         return (int)next.size();
  117.     }
  118. };
  119.  
  120. // ===================================================================
  121. // 2) Basic queries on a SAM
  122. // ===================================================================
  123.  
  124. // 2.1) Count the number of distinct substrings of the original string.
  125. //      Input: sam (built)
  126. //      Returns: long long number of distinct substrings.
  127. //      Time: O(number of states) = O(n).
  128. //      Formula: sum over all states (len[state] - len[link[state]]).
  129. long long countDistinctSubstrings(const SuffixAutomaton& sam) {
  130.     long long ans = 0;
  131.     for (int i = 1; i < sam.size(); ++i) {
  132.         ans += sam.len[i] - sam.len[sam.link[i]];
  133.     }
  134.     return ans;
  135. }
  136.  
  137. // 2.2) Check if a pattern exists as a substring.
  138. //      Input: sam, pattern p
  139. //      Returns: true if p is a substring, false otherwise.
  140. //      Time: O(|p|).
  141. bool substringExists(const SuffixAutomaton& sam, const string& p) {
  142.     int state = 0;
  143.     for (char ch : p) {
  144.         int c = ch - 'a';
  145.         if (sam.next[state][c] == -1) return false;
  146.         state = sam.next[state][c];
  147.     }
  148.     return true;
  149. }
  150.  
  151. // 2.3) Get the occurrence count of a given pattern.
  152. //      Requires pre‑computed occ vector.
  153. //      Input: sam, occ, pattern p
  154. //      Returns: int number of occurrences.
  155. //      Time: O(|p|).
  156. int getOccurrenceCount(const SuffixAutomaton& sam, const vector<int>& occ, const string& p) {
  157.     int state = 0;
  158.     for (char ch : p) {
  159.         int c = ch - 'a';
  160.         if (sam.next[state][c] == -1) return 0;
  161.         state = sam.next[state][c];
  162.     }
  163.     return occ[state];
  164. }
  165.  
  166. // 2.4) Compute occurrence counts for all states.
  167. //      Input: sam
  168. //      Returns: vector<int> occ where occ[state] = number of occurrences.
  169. //      Time: O(number of states).
  170. //      Note: Uses isClone to identify non‑clone states (initially 1).
  171. vector<int> computeOccurrences(const SuffixAutomaton& sam) {
  172.     int n = sam.size();
  173.     vector<int> occ(n, 0);
  174.     for (int i = 1; i < n; ++i) {
  175.         if (!sam.isClone[i]) occ[i] = 1;
  176.     }
  177.     // Order states by length descending.
  178.     vector<int> order(n);
  179.     iota(order.begin(), order.end(), 0);
  180.     sort(order.begin(), order.end(), [&](int a, int b) {
  181.         return sam.len[a] > sam.len[b];
  182.     });
  183.     for (int v : order) {
  184.         if (sam.link[v] != -1) {
  185.             occ[sam.link[v]] += occ[v];
  186.         }
  187.     }
  188.     return occ;
  189. }
  190.  
  191. // ===================================================================
  192. // 3) Advanced queries on a SAM
  193. // ===================================================================
  194.  
  195. // 3.1) Longest common substring between two strings.
  196. //      Input: SAM built from s, second string t
  197. //      Returns: int length of the longest common substring.
  198. //      Time: O(|t|).
  199. int longestCommonSubstring(const SuffixAutomaton& sam, const string& t) {
  200.     int state = 0;
  201.     int curLen = 0;
  202.     int best = 0;
  203.     for (char ch : t) {
  204.         int c = ch - 'a';
  205.         if (sam.next[state][c] != -1) {
  206.             state = sam.next[state][c];
  207.             curLen++;
  208.         } else {
  209.             while (state != -1 && sam.next[state][c] == -1) {
  210.                 state = sam.link[state];
  211.             }
  212.             if (state == -1) {
  213.                 state = 0;
  214.                 curLen = 0;
  215.             } else {
  216.                 curLen = sam.len[state] + 1;
  217.                 state = sam.next[state][c];
  218.             }
  219.         }
  220.         best = max(best, curLen);
  221.     }
  222.     return best;
  223. }
  224.  
  225. // 3.2) Longest repeated substring (at least twice, can overlap).
  226. //      Input: sam, occ (from computeOccurrences)
  227. //      Returns: int length of the longest repeated substring.
  228. //      Time: O(number of states).
  229. int longestRepeatedSubstring(const SuffixAutomaton& sam, const vector<int>& occ) {
  230.     int ans = 0;
  231.     for (int i = 1; i < sam.size(); ++i) {
  232.         if (occ[i] >= 2) {
  233.             ans = max(ans, sam.len[i]);
  234.         }
  235.     }
  236.     return ans;
  237. }
  238.  
  239. // 3.3) Compute first and last occurrence positions for each state.
  240. //      Input: sam
  241. //      Returns: pair<vector<int>, vector<int>> (firstPos, lastPos)
  242. //      Time: O(number of states).
  243. //      Note: Initial positions are set for non‑clone states (len[state]-1),
  244. //            then propagated along suffix links.
  245. pair<vector<int>, vector<int>> computeFirstLastPos(const SuffixAutomaton& sam) {
  246.     int n = sam.size();
  247.     vector<int> firstPos(n, INF);
  248.     vector<int> lastPos(n, -INF);
  249.     for (int i = 1; i < n; ++i) {
  250.         if (!sam.isClone[i]) {
  251.             firstPos[i] = lastPos[i] = sam.len[i] - 1;
  252.         }
  253.     }
  254.     vector<int> order(n);
  255.     iota(order.begin(), order.end(), 0);
  256.     sort(order.begin(), order.end(), [&](int a, int b) {
  257.         return sam.len[a] > sam.len[b];
  258.     });
  259.     for (int v : order) {
  260.         if (sam.link[v] != -1) {
  261.             int p = sam.link[v];
  262.             firstPos[p] = min(firstPos[p], firstPos[v]);
  263.             lastPos[p] = max(lastPos[p], lastPos[v]);
  264.         }
  265.     }
  266.     return {firstPos, lastPos};
  267. }
  268.  
  269. // 3.4) Longest repeated substring that does NOT overlap.
  270. //      Input: sam, firstPos, lastPos (from computeFirstLastPos)
  271. //      Returns: int length of the longest non‑overlapping repeated substring.
  272. //      Time: O(number of states).
  273. int longestNonOverlappingRepeated(const SuffixAutomaton& sam,
  274.                                   const vector<int>& firstPos,
  275.                                   const vector<int>& lastPos) {
  276.     int ans = 0;
  277.     for (int v = 1; v < sam.size(); ++v) {
  278.         if (firstPos[v] + sam.len[v] <= lastPos[v]) {
  279.             ans = max(ans, sam.len[v]);
  280.         }
  281.     }
  282.     return ans;
  283. }
  284.  
  285. // ===================================================================
  286. // 4) Lexicographical queries on a SAM
  287. // ===================================================================
  288.  
  289. // 4.1) Count the number of distinct substrings that start from each state.
  290. //      Input: sam
  291. //      Returns: vector<long long> dp where dp[state] = number of distinct
  292. //               substrings (including the empty string) starting from state.
  293. //      Time: O(number of states + transitions).
  294. vector<long long> computeDP(const SuffixAutomaton& sam) {
  295.     int n = sam.size();
  296.     vector<long long> dp(n, 0);
  297.     vector<int> order(n);
  298.     iota(order.begin(), order.end(), 0);
  299.     sort(order.begin(), order.end(), [&](int a, int b) {
  300.         return sam.len[a] > sam.len[b];
  301.     });
  302.     for (int v : order) {
  303.         dp[v] = 1; // empty string
  304.         for (int c = 0; c < sam.ALPHABET; ++c) {
  305.             if (sam.next[v][c] != -1) {
  306.                 dp[v] += dp[sam.next[v][c]];
  307.             }
  308.         }
  309.     }
  310.     return dp;
  311. }
  312.  
  313. // 4.2) Find the k‑th lexicographically smallest distinct substring (1‑indexed).
  314. //      Input: sam, dp (from computeDP), k (1‑based)
  315. //      Returns: string – the k‑th distinct substring.
  316. //      Time: O(answer length * alphabet).
  317. //      Note: dp includes the empty string; we skip it.
  318. string kthSmallestSubstring(const SuffixAutomaton& sam, const vector<long long>& dp, long long k) {
  319.     string ans;
  320.     int state = 0;
  321.     while (k > 0) {
  322.         for (int c = 0; c < sam.ALPHABET; ++c) {
  323.             if (sam.next[state][c] != -1) {
  324.                 int nxt = sam.next[state][c];
  325.                 if (dp[nxt] >= k) {
  326.                     ans.push_back(char('a' + c));
  327.                     state = nxt;
  328.                     k--; // consumed the empty continuation of this branch
  329.                     break;
  330.                 } else {
  331.                     k -= dp[nxt];
  332.                 }
  333.             }
  334.         }
  335.     }
  336.     return ans;
  337. }
  338.  
  339. // ===================================================================
  340. // 5) Advanced tricks and patterns that appeared in ECPC/ACPC
  341. // ===================================================================
  342.  
  343. // 5.1) Minimum lexicographic rotation of a string.
  344. //      Input: s
  345. //      Returns: string – the smallest rotation.
  346. //      Time: O(n) using SAM on s+s.
  347. //      Note: We greedily follow the smallest transition for n steps.
  348. string minLexicographicRotation(const string& s) {
  349.     string ss = s + s;
  350.     SuffixAutomaton sam(ss);
  351.     string ans;
  352.     int state = 0;
  353.     for (int i = 0; i < (int)s.size(); ++i) {
  354.         for (int c = 0; c < sam.ALPHABET; ++c) {
  355.             if (sam.next[state][c] != -1) {
  356.                 ans.push_back(char('a' + c));
  357.                 state = sam.next[state][c];
  358.                 break;
  359.             }
  360.         }
  361.     }
  362.     return ans;
  363. }
  364.  
  365. // 5.2) Generalized Suffix Automaton (for multiple strings).
  366. //      Input: vector of strings
  367. //      Returns: SuffixAutomaton built from all strings.
  368. //      Time: O(total length * alphabet).
  369. //      Note: Resets `last` to 0 before each string.
  370. SuffixAutomaton buildGeneralizedSAM(const vector<string>& strings) {
  371.     SuffixAutomaton sam;
  372.     sam.next.push_back({});
  373.     for (auto &a : sam.next[0]) a = -1;
  374.     sam.link.push_back(-1);
  375.     sam.len.push_back(0);
  376.     sam.isClone.push_back(false);
  377.     sam.occ.push_back(0);
  378.     sam.firstPos.push_back(INF);
  379.     sam.lastPos.push_back(-INF);
  380.     sam.last = 0;
  381.  
  382.     for (const string& s : strings) {
  383.         sam.last = 0;
  384.         for (char ch : s) {
  385.             sam.extend(ch - 'a');
  386.         }
  387.     }
  388.     return sam;
  389. }
  390.  
  391. // 5.3) Count substrings that appear in at least K of the given strings.
  392. //      Input: vector of strings, K
  393. //      Returns: long long number of distinct substrings appearing in >= K strings.
  394. //      Time: O(total length * alphabet + states * number_of_strings).
  395. //      Note: Uses masks (sets) for each state and propagates along suffix links.
  396. long long countSubstringsInAtLeastK(const vector<string>& strings, int K) {
  397.     SuffixAutomaton sam = buildGeneralizedSAM(strings);
  398.     int n_states = sam.size();
  399.     int m = strings.size();
  400.  
  401.     vector<set<int>> masks(n_states);
  402.     // For each string, mark the states visited by its prefixes.
  403.     for (int idx = 0; idx < m; ++idx) {
  404.         int state = 0;
  405.         for (char ch : strings[idx]) {
  406.             int c = ch - 'a';
  407.             if (sam.next[state][c] == -1) break; // should not happen
  408.             state = sam.next[state][c];
  409.             masks[state].insert(idx);
  410.         }
  411.     }
  412.  
  413.     // Propagate masks along suffix links.
  414.     vector<int> order(n_states);
  415.     iota(order.begin(), order.end(), 0);
  416.     sort(order.begin(), order.end(), [&](int a, int b) {
  417.         return sam.len[a] > sam.len[b];
  418.     });
  419.     for (int v : order) {
  420.         if (sam.link[v] != -1) {
  421.             int p = sam.link[v];
  422.             // Union (small-to-large)
  423.             if (masks[v].size() > masks[p].size()) swap(masks[v], masks[p]);
  424.             for (int x : masks[v]) masks[p].insert(x);
  425.         }
  426.     }
  427.  
  428.     long long ans = 0;
  429.     for (int i = 1; i < n_states; ++i) {
  430.         if ((int)masks[i].size() >= K) {
  431.             ans += sam.len[i] - sam.len[sam.link[i]];
  432.         }
  433.     }
  434.     return ans;
  435. }
  436.  
  437. // ===================================================================
  438. // 6) Helper: topological order of states by length
  439. // ===================================================================
  440.  
  441. // 6.1) Get states sorted by length (descending).
  442. //      Input: sam
  443. //      Returns: vector<int> states in decreasing order of len.
  444. vector<int> getStatesByLengthDesc(const SuffixAutomaton& sam) {
  445.     int n = sam.size();
  446.     vector<int> order(n);
  447.     iota(order.begin(), order.end(), 0);
  448.     sort(order.begin(), order.end(), [&](int a, int b) {
  449.         return sam.len[a] > sam.len[b];
  450.     });
  451.     return order;
  452. }
  453.  
  454. // ===================================================================
  455. // main() with example usage (you can ignore this part)
  456. // ===================================================================
  457.  
  458. int main() {
  459.     ios::sync_with_stdio(false);
  460.     cin.tie(nullptr);
  461.  
  462.     string s = "ababa";
  463.     SuffixAutomaton sam(s);
  464.  
  465.     // Count distinct substrings
  466.     cout << "Distinct substrings: " << countDistinctSubstrings(sam) << '\n'; // 9
  467.  
  468.     // Check existence
  469.     cout << "Contains 'aba'? " << substringExists(sam, "aba") << '\n'; // 1
  470.  
  471.     // Compute occurrences
  472.     vector<int> occ = computeOccurrences(sam);
  473.     cout << "Occurrences of 'ba': " << getOccurrenceCount(sam, occ, "ba") << '\n'; // 2
  474.  
  475.     // Longest common substring with another string
  476.     string t = "bab";
  477.     cout << "LCS length: " << longestCommonSubstring(sam, t) << '\n'; // 2
  478.  
  479.     // Longest repeated substring
  480.     cout << "Longest repeated: " << longestRepeatedSubstring(sam, occ) << '\n'; // 3
  481.  
  482.     // Non‑overlapping repeated
  483.     auto [firstPos, lastPos] = computeFirstLastPos(sam);
  484.     cout << "Longest non‑overlapping repeated: " << longestNonOverlappingRepeated(sam, firstPos, lastPos) << '\n'; // e.g., 1
  485.  
  486.     // k‑th smallest substring
  487.     vector<long long> dp = computeDP(sam);
  488.     cout << "3rd smallest substring: " << kthSmallestSubstring(sam, dp, 3) << '\n'; // "ab"
  489.  
  490.     // Minimum rotation
  491.     string rot = "bca";
  492.     cout << "Min rotation of " << rot << ": " << minLexicographicRotation(rot) << '\n'; // "abc"
  493.  
  494.     // Count substrings in at least K strings
  495.     vector<string> strs = {"ab", "bc", "abc"};
  496.     cout << "Substrings appearing in at least 2 strings: " << countSubstringsInAtLeastK(strs, 2) << '\n'; // e.g., "b", "bc", "c"?
  497.  
  498.     return 0;
  499. }
Success #stdin #stdout 0.01s 5288KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Distinct substrings: 9
Contains 'aba'? 1
Occurrences of 'ba': 2
LCS length: 3
Longest repeated: 3
Longest non‑overlapping repeated: 2
3rd smallest substring: aba
Min rotation of bca: abc
Substrings appearing in at least 2 strings: 2
https://ideone.com/gmL4nW
language:
C++ (gcc 8.3)
created:
1 day ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!