Problem p03355 (C++)

Buggy submission (status: Runtime Error):
// #include <ext/pb_ds/assoc_container.hpp> // Common file
// #include <ext/pb_ds/tree_policy.hpp> // Including
// tree_order_statistics_node_update #include
// <ext/pb_ds/detail/standard_policies.hpp>
#include <algorithm>
#include <assert.h>
#include <bitset>
#include <cmath>
#include <complex>
#include <fstream>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <memory.h>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdio.h>
#include <string>
#include <strings.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define f first
#define s second
#define pb push_back
#define lp(i, a, n) for (int(i) = (a); (i) <= (int)(n); (i)++)
#define lpd(i, n, a) for (int(i) = (n); (i) >= (a); --(i))
#define mp make_pair
#define clr(a, b) memset(a, b, sizeof a)
#define all(v) v.begin(), v.end()
#define mod 1000000007
#define eps 1e-8
#define infi 2e9
#define infll 1e18
#define MX 1000000
#define X real()
#define Y imag()
#define polar(r, t) ((r)*exp(point(0, (t))))
#define length(a) hypot((a).X, (a).Y)
#define angle(a) atan2((a).Y, (a).X)
#define vec(a, b) ((b) - (a))
#define dot(a, b) (conj(a) * (b)).X
#define cross(a, b) (conj(a) * (b)).Y
#define lengthsqr(a) dot(a, a)
#define reflect(V, M) (conj((V) / (M)) * (M))
// #define rotate(a, theta) (a) * (polar(1.0, (theta) ))
#define normalize(a) ((a) / length(a))
#define ccw(a, b, c) cross(vec(a, b), vec(a, c)) > eps
#define cosRule(a, b, c)                                                       \
  (acos(((a) * (a) + (b) * (b) - (c) * (c)) / (2 * (a) * (b))))
#define cosDot(a, b) (acos(dot(a, b) / length(a) / length(b)))
#define EQ(a, b) (fabs((a) - (b)) <= eps) /* equal to */
#define NE(a, b) (fabs((a) - (b)) > eps)  /* not equal to */
#define LT(a, b) ((a) < (b)-eps)          /* less than */
#define GT(a, b) ((a) > (b) + eps)        /* greater than */
#define LE(a, b) ((a) <= (b) + eps)       /* less than or equal to */
#define GE(a, b) ((a) >= (b)-eps)         /* greater than or equal to */
#define mod1 100050001
#define mod2 100030001
#define base1 37
#define base2 31
#define PI acos(-1)
#define sz(a) (int)(a).size()

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<double, double> pdd;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<vi> vvi;
typedef set<int> si;
typedef map<int, int> mii;
typedef complex<long double> point;
typedef pair<point, point> line;
typedef pair<double, point> Circle;
typedef long double ld;

const int N = 5005, ALPHA = 28;

struct trie {
  map<int, trie *> childs;
  int siz = 0;

  trie() { siz = 0; };

  int insert(char *S) {
    siz = max(siz, 1);
    if (*S == '\0')
      return 0;

    int p = *S - 'a', ret = 0;
    if (!childs.count(p))
      childs[p] = new trie(), ++ret;

    ret += childs[p]->insert(S + 1);
    siz += ret;
    return ret;
  }

  void calc_siz() {
    siz = 1;
    lp(i, 0, ALPHA - 1) if (childs[i]) {
      childs[i]->calc_siz();
      siz += childs[i]->siz;
    }
  }

  void search(string &ans, int k) {

    if (!k)
      return;

    lp(i, 0, ALPHA - 1) if (childs[i]) {
      if (childs[i]->siz < k)
        k -= childs[i]->siz;
      else {
        ans += i + 'a', childs[i]->search(ans, k - 1);
        return;
      }
    }
  }
};

unsigned char mask[] = {0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01};
#define tget(i) ((t[(i) / 8] & mask[(i) % 8]) ? 1 : 0)
#define tset(i, b)                                                             \
  t[(i) / 8] =                                                                 \
      (b) ? (mask[(i) % 8] | t[(i) / 8]) : ((~mask[(i) % 8]) & t[(i) / 8])
#define chr(i) (cs == sizeof(int) ? ((int *)s)[i] : ((unsigned char *)s)[i])
#define isLMS(i) (i > 0 && tget(i) && !tget(i - 1))

// find the start or end of each bucket
void getBuckets(unsigned char *s, int *bkt, int n, int K, int cs, bool end) {
  int i, sum = 0;
  for (i = 0; i <= K; i++)
    bkt[i] = 0; // clear all buckets
  for (i = 0; i < n; i++)
    bkt[chr(i)]++; // compute the size of each bucket
  for (i = 0; i <= K; i++) {
    sum += bkt[i];
    bkt[i] = end ? sum : sum - bkt[i];
  }
}
// compute SAl
void induceSAl(unsigned char *t, int *SA, unsigned char *s, int *bkt, int n,
               int K, int cs, bool end) {
  int i, j;
  getBuckets(s, bkt, n, K, cs, end); // find starts of buckets
  for (i = 0; i < n; i++) {
    j = SA[i] - 1;
    if (j >= 0 && !tget(j))
      SA[bkt[chr(j)]++] = j;
  }
}
// compute SAs
void induceSAs(unsigned char *t, int *SA, unsigned char *s, int *bkt, int n,
               int K, int cs, bool end) {
  int i, j;
  getBuckets(s, bkt, n, K, cs, end); // find ends of buckets
  for (i = n - 1; i >= 0; i--) {
    j = SA[i] - 1;
    if (j >= 0 && tget(j))
      SA[--bkt[chr(j)]] = j;
  }
}

// find the suffix array SA of s[0..n-1] in {1..K}^n
// require s[n-1]=0 (the sentinel!), n>=2
// use a working space (excluding s and SA) of at most 2.25n+O(1) for a constant
// alphabet
void SA_IS(unsigned char *s, int *SA, int n, int K, int cs) {
  int i, j;
  unsigned char *t =
      (unsigned char *)malloc(n / 8 + 1); // LS-type array in bits
  // Classify the type of each character
  tset(n - 2, 0);
  tset(n - 1, 1); // the sentinel must be in s1, important!!!
  for (i = n - 3; i >= 0; i--)
    tset(i, (chr(i) < chr(i + 1) || (chr(i) == chr(i + 1) && tget(i + 1) == 1))
                ? 1
                : 0);
  // stage 1: reduce the problem by at least 1/2
  // sort all the S-substrings
  int *bkt = (int *)malloc(sizeof(int) * (K + 1)); // bucket array
  getBuckets(s, bkt, n, K, cs, true);              // find ends of buckets
  for (i = 0; i < n; i++)
    SA[i] = -1;
  for (i = 1; i < n; i++)
    if (isLMS(i))
      SA[--bkt[chr(i)]] = i;
  induceSAl(t, SA, s, bkt, n, K, cs, false);
  induceSAs(t, SA, s, bkt, n, K, cs, true);
  free(bkt);
  // compact all the sorted substrings into the first n1 items of SA
  // 2*n1 must be not larger than n (proveable)
  int n1 = 0;
  for (i = 0; i < n; i++)
    if (isLMS(SA[i]))
      SA[n1++] = SA[i];
  // find the lexicographic names of all substrings
  for (i = n1; i < n; i++)
    SA[i] = -1; // init the name array buffer
  int name = 0, prev = -1;
  for (i = 0; i < n1; i++) {
    int pos = SA[i];
    bool diff = false;
    for (int d = 0; d < n; d++)
      if (prev == -1 || chr(pos + d) != chr(prev + d) ||
          tget(pos + d) != tget(prev + d)) {
        diff = true;
        break;
      } else if (d > 0 && (isLMS(pos + d) || isLMS(prev + d)))
        break;
    if (diff) {
      name++;
      prev = pos;
    }
    pos = (pos % 2 == 0) ? pos / 2 : (pos - 1) / 2;
    SA[n1 + pos] = name - 1;
  }
  for (i = n - 1, j = n - 1; i >= n1; i--)
    if (SA[i] >= 0)
      SA[j--] = SA[i];
  // stage 2: solve the reduced problem
  // recurse if names are not yet unique
  int *SA1 = SA, *s1 = SA + n - n1;
  if (name < n1)
    SA_IS((unsigned char *)s1, SA1, n1, name - 1, sizeof(int));
  else
    // generate the suffix array of s1 directly
    for (i = 0; i < n1; i++)
      SA1[s1[i]] = i;
  // stage 3: induce the result for the original problem
  bkt = (int *)malloc(sizeof(int) * (K + 1)); // bucket array
  // put all left-most S characters into their buckets
  getBuckets(s, bkt, n, K, cs, true); // find ends of buckets
  for (i = 1, j = 0; i < n; i++)
    if (isLMS(i))
      s1[j++] = i; // get p1
  for (i = 0; i < n1; i++)
    SA1[i] = s1[SA1[i]]; // get index in s
  for (i = n1; i < n; i++)
    SA[i] = -1; // init SA[n1..n-1]
  for (i = n1 - 1; i >= 0; i--) {
    j = SA[i];
    SA[i] = -1;
    SA[--bkt[chr(j)]] = j;
  }
  induceSAl(t, SA, s, bkt, n, K, cs, false);
  induceSAs(t, SA, s, bkt, n, K, cs, true);
  free(bkt);
  free(t);
}

int sa[N];
int lcp[N];
int Rank[N];
unsigned char *s;

void calc_lcp() {
  // compute the rank of each suffix
  for (int i = 0; i - 1 < 0 || s[i - 1]; i++)
    Rank[sa[i]] = i;

  int c = 0; // the length of lcp between i and j
  for (int i = 0; i - 1 < 0 || s[i - 1]; i++) {
    if (Rank[i]) { // if i is not the first suffix in the sorted array
      int j = sa[Rank[i] - 1]; // find the element before i and name it j
      while (s[i + c] == s[j + c])
        c++; // count the number of shared chars
    }
    lcp[Rank[i]] = c; // store the result in lcp array
    if (c)
      c--; // Decrement c by one because length of lcp of i+1 is c-1
  }
}

//
// int main() {
//    string str = "abcab";
//    n = sz(str);
//    s = (unsigned char*) str.c_str();
//    SA_IS(s, sa, n + 1, 256, 1);
//    calc_lcp();
//
//
//
//
//}

int n, k;
char str[N];
trie tr;
vi ind[ALPHA];

void print(string str) {
  for (int i = 0; i < n; i++) {
    cout << str.substr(sa[i + 1]);
    if (i < n - 1)
      cout << " " << lcp[i + 1];
    cout << endl;
  }
}

int main() {

  scanf("%s%d", str, &k);
  s = (unsigned char *)string(str).c_str();
  n = (int)strlen(str);

  SA_IS(s, sa, n + 1, 256, 1);
  calc_lcp();

  lp(i, 0, min(2 * k, n)) { tr.insert(str + sa[i]); }

  string ans;
  tr.search(ans, k);

  printf("%s\n", ans.c_str());
  return 0;
}

Fixed submission:
// #include <ext/pb_ds/assoc_container.hpp> // Common file
// #include <ext/pb_ds/tree_policy.hpp> // Including
// tree_order_statistics_node_update #include
// <ext/pb_ds/detail/standard_policies.hpp>
#include <algorithm>
#include <assert.h>
#include <bitset>
#include <cmath>
#include <complex>
#include <fstream>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <math.h>
#include <memory.h>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <stdio.h>
#include <string>
#include <strings.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>

#define f first
#define s second
#define pb push_back
#define lp(i, a, n) for (int(i) = (a); (i) <= (int)(n); (i)++)
#define lpd(i, n, a) for (int(i) = (n); (i) >= (a); --(i))
#define mp make_pair
#define clr(a, b) memset(a, b, sizeof a)
#define all(v) v.begin(), v.end()
#define mod 1000000007
#define eps 1e-8
#define infi 2e9
#define infll 1e18
#define MX 1000000
#define X real()
#define Y imag()
#define polar(r, t) ((r)*exp(point(0, (t))))
#define length(a) hypot((a).X, (a).Y)
#define angle(a) atan2((a).Y, (a).X)
#define vec(a, b) ((b) - (a))
#define dot(a, b) (conj(a) * (b)).X
#define cross(a, b) (conj(a) * (b)).Y
#define lengthsqr(a) dot(a, a)
#define reflect(V, M) (conj((V) / (M)) * (M))
// #define rotate(a, theta) (a) * (polar(1.0, (theta) ))
#define normalize(a) ((a) / length(a))
#define ccw(a, b, c) cross(vec(a, b), vec(a, c)) > eps
#define cosRule(a, b, c)                                                       \
  (acos(((a) * (a) + (b) * (b) - (c) * (c)) / (2 * (a) * (b))))
#define cosDot(a, b) (acos(dot(a, b) / length(a) / length(b)))
#define EQ(a, b) (fabs((a) - (b)) <= eps) /* equal to */
#define NE(a, b) (fabs((a) - (b)) > eps)  /* not equal to */
#define LT(a, b) ((a) < (b)-eps)          /* less than */
#define GT(a, b) ((a) > (b) + eps)        /* greater than */
#define LE(a, b) ((a) <= (b) + eps)       /* less than or equal to */
#define GE(a, b) ((a) >= (b)-eps)         /* greater than or equal to */
#define mod1 100050001
#define mod2 100030001
#define base1 37
#define base2 31
#define PI acos(-1)
#define sz(a) (int)(a).size()

using namespace std;
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<double, double> pdd;
typedef pair<ll, ll> pll;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef vector<vi> vvi;
typedef set<int> si;
typedef map<int, int> mii;
typedef complex<long double> point;
typedef pair<point, point> line;
typedef pair<double, point> Circle;
typedef long double ld;

const int N = 5005, ALPHA = 28;

struct trie {
  map<int, trie *> childs;
  int siz = 0;

  trie() { siz = 0; };

  int insert(char *S) {
    siz = max(siz, 1);
    if (*S == '\0')
      return 0;

    int p = *S - 'a', ret = 0;
    if (!childs.count(p))
      childs[p] = new trie(), ++ret;

    ret += childs[p]->insert(S + 1);
    siz += ret;
    return ret;
  }

  void calc_siz() {
    siz = 1;
    lp(i, 0, ALPHA - 1) if (childs[i]) {
      childs[i]->calc_siz();
      siz += childs[i]->siz;
    }
  }

  void search(string &ans, int k) {

    if (!k)
      return;

    lp(i, 0, ALPHA - 1) if (childs[i]) {
      if (childs[i]->siz < k)
        k -= childs[i]->siz;
      else {
        ans += i + 'a', childs[i]->search(ans, k - 1);
        return;
      }
    }
  }
};

unsigned char mask[] = {0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01};
#define tget(i) ((t[(i) / 8] & mask[(i) % 8]) ? 1 : 0)
#define tset(i, b)                                                             \
  t[(i) / 8] =                                                                 \
      (b) ? (mask[(i) % 8] | t[(i) / 8]) : ((~mask[(i) % 8]) & t[(i) / 8])
#define chr(i) (cs == sizeof(int) ? ((int *)s)[i] : ((unsigned char *)s)[i])
#define isLMS(i) (i > 0 && tget(i) && !tget(i - 1))

// find the start or end of each bucket
void getBuckets(unsigned char *s, int *bkt, int n, int K, int cs, bool end) {
  int i, sum = 0;
  for (i = 0; i <= K; i++)
    bkt[i] = 0; // clear all buckets
  for (i = 0; i < n; i++)
    bkt[chr(i)]++; // compute the size of each bucket
  for (i = 0; i <= K; i++) {
    sum += bkt[i];
    bkt[i] = end ? sum : sum - bkt[i];
  }
}
// compute SAl
void induceSAl(unsigned char *t, int *SA, unsigned char *s, int *bkt, int n,
               int K, int cs, bool end) {
  int i, j;
  getBuckets(s, bkt, n, K, cs, end); // find starts of buckets
  for (i = 0; i < n; i++) {
    j = SA[i] - 1;
    if (j >= 0 && !tget(j))
      SA[bkt[chr(j)]++] = j;
  }
}
// compute SAs
void induceSAs(unsigned char *t, int *SA, unsigned char *s, int *bkt, int n,
               int K, int cs, bool end) {
  int i, j;
  getBuckets(s, bkt, n, K, cs, end); // find ends of buckets
  for (i = n - 1; i >= 0; i--) {
    j = SA[i] - 1;
    if (j >= 0 && tget(j))
      SA[--bkt[chr(j)]] = j;
  }
}

// find the suffix array SA of s[0..n-1] in {1..K}^n
// require s[n-1]=0 (the sentinel!), n>=2
// use a working space (excluding s and SA) of at most 2.25n+O(1) for a constant
// alphabet
void SA_IS(unsigned char *s, int *SA, int n, int K, int cs) {
  int i, j;
  unsigned char *t =
      (unsigned char *)malloc(n / 8 + 1); // LS-type array in bits
  // Classify the type of each character
  tset(n - 2, 0);
  tset(n - 1, 1); // the sentinel must be in s1, important!!!
  for (i = n - 3; i >= 0; i--)
    tset(i, (chr(i) < chr(i + 1) || (chr(i) == chr(i + 1) && tget(i + 1) == 1))
                ? 1
                : 0);
  // stage 1: reduce the problem by at least 1/2
  // sort all the S-substrings
  int *bkt = (int *)malloc(sizeof(int) * (K + 1)); // bucket array
  getBuckets(s, bkt, n, K, cs, true);              // find ends of buckets
  for (i = 0; i < n; i++)
    SA[i] = -1;
  for (i = 1; i < n; i++)
    if (isLMS(i))
      SA[--bkt[chr(i)]] = i;
  induceSAl(t, SA, s, bkt, n, K, cs, false);
  induceSAs(t, SA, s, bkt, n, K, cs, true);
  free(bkt);
  // compact all the sorted substrings into the first n1 items of SA
  // 2*n1 must be not larger than n (proveable)
  int n1 = 0;
  for (i = 0; i < n; i++)
    if (isLMS(SA[i]))
      SA[n1++] = SA[i];
  // find the lexicographic names of all substrings
  for (i = n1; i < n; i++)
    SA[i] = -1; // init the name array buffer
  int name = 0, prev = -1;
  for (i = 0; i < n1; i++) {
    int pos = SA[i];
    bool diff = false;
    for (int d = 0; d < n; d++)
      if (prev == -1 || chr(pos + d) != chr(prev + d) ||
          tget(pos + d) != tget(prev + d)) {
        diff = true;
        break;
      } else if (d > 0 && (isLMS(pos + d) || isLMS(prev + d)))
        break;
    if (diff) {
      name++;
      prev = pos;
    }
    pos = (pos % 2 == 0) ? pos / 2 : (pos - 1) / 2;
    SA[n1 + pos] = name - 1;
  }
  for (i = n - 1, j = n - 1; i >= n1; i--)
    if (SA[i] >= 0)
      SA[j--] = SA[i];
  // stage 2: solve the reduced problem
  // recurse if names are not yet unique
  int *SA1 = SA, *s1 = SA + n - n1;
  if (name < n1)
    SA_IS((unsigned char *)s1, SA1, n1, name - 1, sizeof(int));
  else
    // generate the suffix array of s1 directly
    for (i = 0; i < n1; i++)
      SA1[s1[i]] = i;
  // stage 3: induce the result for the original problem
  bkt = (int *)malloc(sizeof(int) * (K + 1)); // bucket array
  // put all left-most S characters into their buckets
  getBuckets(s, bkt, n, K, cs, true); // find ends of buckets
  for (i = 1, j = 0; i < n; i++)
    if (isLMS(i))
      s1[j++] = i; // get p1
  for (i = 0; i < n1; i++)
    SA1[i] = s1[SA1[i]]; // get index in s
  for (i = n1; i < n; i++)
    SA[i] = -1; // init SA[n1..n-1]
  for (i = n1 - 1; i >= 0; i--) {
    j = SA[i];
    SA[i] = -1;
    SA[--bkt[chr(j)]] = j;
  }
  induceSAl(t, SA, s, bkt, n, K, cs, false);
  induceSAs(t, SA, s, bkt, n, K, cs, true);
  free(bkt);
  free(t);
}

int sa[N];
int lcp[N];
int Rank[N];
unsigned char *s;

void calc_lcp() {
  // compute the rank of each suffix
  for (int i = 0; i - 1 < 0 || s[i - 1]; i++)
    Rank[sa[i]] = i;

  int c = 0; // the length of lcp between i and j
  for (int i = 0; i - 1 < 0 || s[i - 1]; i++) {
    if (Rank[i]) { // if i is not the first suffix in the sorted array
      int j = sa[Rank[i] - 1]; // find the element before i and name it j
      while (s[i + c] == s[j + c])
        c++; // count the number of shared chars
    }
    lcp[Rank[i]] = c; // store the result in lcp array
    if (c)
      c--; // Decrement c by one because length of lcp of i+1 is c-1
  }
}

//
// int main() {
//    string str = "abcab";
//    n = sz(str);
//    s = (unsigned char*) str.c_str();
//    SA_IS(s, sa, n + 1, 256, 1);
//    calc_lcp();
//
//
//
//
//}

int n, k;
char str[N];
trie tr;
vi ind[ALPHA];

void print(string str) {
  for (int i = 0; i < n; i++) {
    cout << str.substr(sa[i + 1]);
    if (i < n - 1)
      cout << " " << lcp[i + 1];
    cout << endl;
  }
}

int main() {

  scanf("%s%d", str, &k);
  s = (unsigned char *)string(str).c_str();
  n = (int)strlen(str);
  //
  //    SA_IS(s, sa, n + 1, 256, 1);
  //    calc_lcp();

  priority_queue<string, vector<string>, greater<string>> q;
  lp(i, 0, n - 1) q.push(str + i);

  string p, ans;

  while (k) {

    string cur = q.top();
    q.pop();

    ans.clear();
    int i = 0;
    while (i < sz(p) && i < sz(cur) && cur[i] == p[i])
      ans += p[i++];

    while (k && i < sz(cur)) {
      ans += cur[i++];
      --k;
    }

    p = cur;
  }

  printf("%s\n", ans.c_str());
  return 0;
}