fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. static const int MAXLEN = 15;
  5. static const int MAXSUM = 9 * MAXLEN; // 135
  6.  
  7. static bool isPrimeSum[MAXSUM + 1];
  8. static unsigned long long ways[MAXLEN + 1][MAXSUM + 1];
  9. static unsigned long long cntGoodSuffix[MAXLEN + 1][MAXSUM + 1];
  10. static unsigned long long totalGoodLen[MAXLEN + 1];
  11.  
  12. static void precompute_global() {
  13.     // Sieve primes up to 135
  14.     vector<bool> isPrime(MAXSUM + 1, true);
  15.     isPrime[0] = isPrime[1] = false;
  16.     for (int p = 2; p * p <= MAXSUM; ++p) {
  17.         if (isPrime[p]) {
  18.             for (int q = p * p; q <= MAXSUM; q += p) isPrime[q] = false;
  19.         }
  20.     }
  21.     for (int i = 0; i <= MAXSUM; ++i) isPrimeSum[i] = isPrime[i];
  22.  
  23.     // ways[len][sum]
  24.     memset(ways, 0, sizeof(ways));
  25.     ways[0][0] = 1;
  26.     for (int len = 1; len <= MAXLEN; ++len) {
  27.         for (int s = 0; s <= 9 * len; ++s) {
  28.             unsigned long long v = 0;
  29.             for (int d = 0; d <= 9; ++d) {
  30.                 if (s - d >= 0) v += ways[len - 1][s - d];
  31.             }
  32.             ways[len][s] = v;
  33.         }
  34.     }
  35.  
  36.     // cntGoodSuffix[rem][sumPrefix]
  37.     for (int rem = 0; rem <= MAXLEN; ++rem) {
  38.         for (int sumPrefix = 0; sumPrefix <= MAXSUM; ++sumPrefix) {
  39.             unsigned long long total = 0;
  40.             for (int t = 0; t <= 9 * rem; ++t) {
  41.                 int totSum = sumPrefix + t;
  42.                 if (totSum <= MAXSUM && isPrimeSum[totSum]) {
  43.                     total += ways[rem][t];
  44.                 }
  45.             }
  46.             cntGoodSuffix[rem][sumPrefix] = total;
  47.         }
  48.     }
  49.  
  50.     // totalGoodLen[L]: count of good L-digit numbers (leading digit != 0)
  51.     for (int L = 1; L <= MAXLEN; ++L) {
  52.         if (L == 1) {
  53.             unsigned long long c = 0;
  54.             for (int d = 1; d <= 9; ++d) if (isPrimeSum[d]) ++c;
  55.             totalGoodLen[L] = c;
  56.         } else {
  57.             unsigned long long c = 0;
  58.             for (int first = 1; first <= 9; ++first) {
  59.                 c += cntGoodSuffix[L - 1][first];
  60.             }
  61.             totalGoodLen[L] = c;
  62.         }
  63.     }
  64. }
  65.  
  66. struct Solver {
  67.     string pat;
  68.     int m;
  69.  
  70.     // KMP transitions: for state 0..m-1 and digit 0..9 -> next state (0..m)
  71.     vector<array<uint8_t, 10>> kmpTrans;
  72.  
  73.     // monoid states for prefixes (merged):
  74.     // For each monoidId and each startState (0..m-1):
  75.     // - monMatch: 1 if the pattern appears while reading this prefix
  76.     // - monEnd  : resulting KMP state after reading the prefix, assuming no match inside
  77.     vector<vector<uint8_t>> monEnd;
  78.     vector<vector<uint8_t>> monMatch;
  79.     vector<array<int, 10>> monNext;
  80.  
  81.     // DP memo: key -> encoded result
  82.     // result encoding: high bit (0x8000) = match, else value = endState
  83.     unordered_map<uint64_t, uint16_t> dp;
  84.  
  85.     static constexpr uint16_t MATCH_CODE = 0x8000;
  86.  
  87.     explicit Solver(string p) : pat(std::move(p)), m((int)pat.size()) {}
  88.  
  89.     void buildKMP() {
  90.         vector<int> pi(m, 0);
  91.         for (int i = 1, j = 0; i < m; ++i) {
  92.             while (j > 0 && pat[i] != pat[j]) j = pi[j - 1];
  93.             if (pat[i] == pat[j]) ++j;
  94.             pi[i] = j;
  95.         }
  96.  
  97.         kmpTrans.assign(m, array<uint8_t, 10>{});
  98.         for (int st = 0; st < m; ++st) {
  99.             for (int d = 0; d <= 9; ++d) {
  100.                 char c = char('0' + d);
  101.                 int j = st;
  102.                 while (j > 0 && pat[j] != c) j = pi[j - 1];
  103.                 if (pat[j] == c) ++j;
  104.                 kmpTrans[st][d] = (uint8_t)j; // can be m
  105.             }
  106.         }
  107.     }
  108.  
  109.     void buildMonoidStates() {
  110.         monEnd.clear();
  111.         monMatch.clear();
  112.         monNext.clear();
  113.  
  114.         vector<uint8_t> end0(m), match0(m, 0);
  115.         for (int s = 0; s < m; ++s) end0[s] = (uint8_t)s;
  116.  
  117.         auto makeKey = [&](const vector<uint8_t>& en, const vector<uint8_t>& ma) -> string {
  118.             string key;
  119.             key.resize((size_t)2 * m);
  120.             for (int i = 0; i < m; ++i) key[i] = (char)en[i];
  121.             for (int i = 0; i < m; ++i) key[m + i] = (char)ma[i];
  122.             return key;
  123.         };
  124.  
  125.         unordered_map<string, int> idMap;
  126.         idMap.reserve(4096);
  127.         idMap.max_load_factor(0.7f);
  128.  
  129.         string key0 = makeKey(end0, match0);
  130.         idMap.emplace(key0, 0);
  131.  
  132.         monEnd.push_back(end0);
  133.         monMatch.push_back(match0);
  134.         monNext.push_back(array<int, 10>{});
  135.         for (int d = 0; d < 10; ++d) monNext[0][d] = -1;
  136.  
  137.         queue<pair<int, int>> q; // (id, length)
  138.         q.push({0, 0});
  139.  
  140.         while (!q.empty()) {
  141.             auto [sid, len] = q.front();
  142.             q.pop();
  143.             if (len == MAXLEN) continue;
  144.  
  145.             for (int d = 0; d <= 9; ++d) {
  146.                 vector<uint8_t> enNew(m, 0), maNew(m, 0);
  147.                 const auto& enOld = monEnd[sid];
  148.                 const auto& maOld = monMatch[sid];
  149.  
  150.                 for (int s = 0; s < m; ++s) {
  151.                     if (maOld[s]) {
  152.                         maNew[s] = 1;
  153.                         enNew[s] = 0;
  154.                     } else {
  155.                         int cur = enOld[s];      // < m
  156.                         int ns = kmpTrans[cur][d];
  157.                         if (ns == m) {
  158.                             maNew[s] = 1;
  159.                             enNew[s] = 0;
  160.                         } else {
  161.                             maNew[s] = 0;
  162.                             enNew[s] = (uint8_t)ns;
  163.                         }
  164.                     }
  165.                 }
  166.  
  167.                 string key = makeKey(enNew, maNew);
  168.                 auto it = idMap.find(key);
  169.                 int nid;
  170.                 if (it == idMap.end()) {
  171.                     nid = (int)monEnd.size();
  172.                     idMap.emplace(std::move(key), nid);
  173.                     monEnd.push_back(std::move(enNew));
  174.                     monMatch.push_back(std::move(maNew));
  175.                     monNext.push_back(array<int, 10>{});
  176.                     for (int dd = 0; dd < 10; ++dd) monNext[nid][dd] = -1;
  177.                     q.push({nid, len + 1});
  178.                 } else {
  179.                     nid = it->second;
  180.                 }
  181.                 monNext[sid][d] = nid;
  182.             }
  183.         }
  184.     }
  185.  
  186.     static inline uint64_t packKey(int rem, int sum, int monoidId, int startState, bool firstFlag) {
  187.         // bits: rem(4) sum(8) monoid(13) start(7) first(1)
  188.         return (uint64_t)rem
  189.              | ((uint64_t)sum << 4)
  190.              | ((uint64_t)monoidId << 12)
  191.              | ((uint64_t)startState << 25)
  192.              | ((uint64_t)(firstFlag ? 1 : 0) << 32);
  193.     }
  194.  
  195.     inline uint16_t solve(int rem, int sum, int monoidId, int startState, bool firstFlag) {
  196.         if (startState == m) return MATCH_CODE;
  197.  
  198.         uint64_t key = packKey(rem, sum, monoidId, startState, firstFlag);
  199.         auto it = dp.find(key);
  200.         if (it != dp.end()) return it->second;
  201.  
  202.         uint16_t ans = 0;
  203.  
  204.         if (rem == 0) {
  205.             if (!isPrimeSum[sum]) {
  206.                 ans = (uint16_t)startState;
  207.             } else if (monMatch[monoidId][startState]) {
  208.                 ans = MATCH_CODE;
  209.             } else {
  210.                 ans = (uint16_t)monEnd[monoidId][startState];
  211.             }
  212.         } else {
  213.             if (cntGoodSuffix[rem][sum] == 0ULL) {
  214.                 // no good numbers at all in this subtree => empty output
  215.                 ans = (uint16_t)startState;
  216.             } else {
  217.                 int cur = startState;
  218.                 int d0 = firstFlag ? 1 : 0;
  219.                 bool found = false;
  220.                 for (int d = d0; d <= 9; ++d) {
  221.                     int newSum = sum + d;
  222.                     if (newSum > MAXSUM) continue;
  223.                     if (cntGoodSuffix[rem - 1][newSum] == 0ULL) continue;
  224.  
  225.                     int childId = monNext[monoidId][d];
  226.                     uint16_t childRes = solve(rem - 1, newSum, childId, cur, false);
  227.                     if (childRes & MATCH_CODE) {
  228.                         ans = MATCH_CODE;
  229.                         found = true;
  230.                         break;
  231.                     }
  232.                     cur = (int)childRes; // endState
  233.                 }
  234.                 if (!found) ans = (uint16_t)cur;
  235.             }
  236.         }
  237.  
  238.         dp.emplace(key, ans);
  239.         return ans;
  240.     }
  241.  
  242.     unsigned long long findEndPosInLengthGroup(
  243.         int L, int rem, int sum, int monoidId, int startState, bool firstFlag,
  244.         unsigned long long offset, char digits[], int depth
  245.     ) {
  246.         if (rem == 0) {
  247.             int cur = startState;
  248.             for (int i = 0; i < L; ++i) {
  249.                 int d = digits[i] - '0';
  250.                 cur = kmpTrans[cur][d];
  251.                 if (cur == m) {
  252.                     return offset + (unsigned long long)(i + 1); // 1-indexed within the number
  253.                 }
  254.             }
  255.             // should not happen
  256.             return offset + (unsigned long long)L;
  257.         }
  258.  
  259.         int cur = startState;
  260.         int d0 = firstFlag ? 1 : 0;
  261.         for (int d = d0; d <= 9; ++d) {
  262.             int newSum = sum + d;
  263.             if (newSum > MAXSUM) continue;
  264.             unsigned long long countChild = cntGoodSuffix[rem - 1][newSum];
  265.             if (countChild == 0ULL) continue;
  266.  
  267.             int childId = monNext[monoidId][d];
  268.             uint16_t childRes = solve(rem - 1, newSum, childId, cur, false);
  269.  
  270.             if (childRes & MATCH_CODE) {
  271.                 digits[depth] = char('0' + d);
  272.                 return findEndPosInLengthGroup(L, rem - 1, newSum, childId, cur, false, offset, digits, depth + 1);
  273.             } else {
  274.                 offset += countChild * (unsigned long long)L;
  275.                 cur = (int)childRes;
  276.             }
  277.         }
  278.  
  279.         // should not happen
  280.         return offset;
  281.     }
  282.  
  283.     unsigned long long solveOne() {
  284.         buildKMP();
  285.         buildMonoidStates();
  286.  
  287.         dp.clear();
  288.         dp.reserve(200000);
  289.         dp.max_load_factor(0.7f);
  290.  
  291.         unsigned long long posBefore = 0; // digits before current length group
  292.         int curState = 0;
  293.  
  294.         for (int L = 1; L <= MAXLEN; ++L) {
  295.             uint16_t res = solve(L, 0, 0, curState, true);
  296.             if (!(res & MATCH_CODE)) {
  297.                 posBefore += totalGoodLen[L] * (unsigned long long)L;
  298.                 curState = (int)res;
  299.             } else {
  300.                 char digits[MAXLEN];
  301.                 unsigned long long endPosInGroup =
  302.                     findEndPosInLengthGroup(L, L, 0, 0, curState, true, 0ULL, digits, 0);
  303.  
  304.                 unsigned long long globalEnd = posBefore + endPosInGroup;
  305.                 unsigned long long startPos = globalEnd - (unsigned long long)m + 1ULL;
  306.                 return startPos;
  307.             }
  308.         }
  309.         // Input guarantees pattern exists
  310.         return 0;
  311.     }
  312. };
  313.  
  314. int main() {
  315.     ios::sync_with_stdio(false);
  316.     cin.tie(nullptr);
  317.  
  318.     precompute_global();
  319.  
  320.     int T;
  321.     cin >> T;
  322.     while (T--) {
  323.         string x;
  324.         cin >> x;
  325.         Solver solver(x);
  326.         cout << solver.solveOne() << "\n";
  327.     }
  328.     return 0;
  329. }
  330.  
Success #stdin #stdout 0.01s 5284KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/P4GclO
language:
C++14 (gcc 8.3)
created:
4 days 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!