Problem p03904 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define ii pair<int, int>
#define app push_back
#define all(a) a.begin(), a.end()
#define bp __builtin_popcountll
#define ll long long
#define mp make_pair
#define f first
#define s second
#define Time (double)clock() / CLOCKS_PER_SEC
#define debug(x) std::cout << #x << ": " << x << '\n';

// It's strictly forbidden to use buildSA() and not
// to add a fictitious integer, which should be smaller
// than all other integers representing a string, to the end!!!
struct SuffixArray {
  static const int SZ = 3e5;

  int c[SZ];
  int cnt[SZ];
  int p[SZ];
  int pn[SZ];
  int cn[SZ];

  vector<int> buildSA(const vector<int> &s) {
    int n = s.size();
    int alpha = (*max_element(s.begin(), s.end())) + 1;
    memset(cnt, 0, alpha * sizeof(int));
    for (int c : s)
      ++cnt[c];
    for (int i = 1; i < alpha; ++i)
      cnt[i] += cnt[i - 1];
    for (int i = 0; i < n; ++i)
      p[--cnt[s[i]]] = i;
    c[p[0]] = 0;
    int cs = 1;
    for (int i = 1; i < n; ++i) {
      if (s[p[i]] != s[p[i - 1]])
        ++cs;
      c[p[i]] = cs - 1;
    }

    for (int h = 0; (1 << h) < n; ++h) {
      for (int i = 0; i < n; ++i) {
        pn[i] = p[i] - (1 << h);
        if (pn[i] < 0)
          pn[i] += n;
      }
      memset(cnt, 0, cs * sizeof(int));
      for (int i = 0; i < n; ++i)
        ++cnt[c[pn[i]]];
      for (int i = 1; i < cs; ++i)
        cnt[i] += cnt[i - 1];
      for (int i = n - 1; i >= 0; --i)
        p[--cnt[c[pn[i]]]] = pn[i];
      cn[p[0]] = 0;
      cs = 1;
      for (int i = 1; i < n; ++i) {
        int mid1 = (p[i] + (1 << h)) % n, mid2 = (p[i - 1] + (1 << h)) % n;
        if (c[p[i]] != c[p[i - 1]] || c[mid1] != c[mid2])
          ++cs;
        cn[p[i]] = cs - 1;
      }
      memcpy(c, cn, n * sizeof(int));
    }

    vector<int> result(p, p + n);
    return result;
  }

  // suf = sa from func above
  vector<int> buildLCP(const vector<int> &s, const vector<int> &suf) const {
    int n = s.size();
    vector<int> rsuf(n);
    vector<int> lcp(n);
    for (ll i = 0; i < n; i++) {
      rsuf[suf[i]] = i;
    }

    int k = 0;
    for (int i = 0; i < n; ++i) {
      if (k > 0)
        --k;
      if (rsuf[i] == n - 1) {
        lcp[n - 1] = -1;
        k = 0;
        continue;
      } else {
        int j = suf[rsuf[i] + 1];
        while (max(i + k, j + k) < n && s[i + k] == s[j + k])
          ++k;
        lcp[rsuf[i]] = k;
      }
    }

    return lcp;
  }
} SA;

const int N = 1e5 + 7, LG = 20;
int dp[N];
int sp[LG][N];

int pw[N];

signed main() {
#ifdef HOME
  freopen("input.txt", "r", stdin);
#else
#define endl '\n'
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cout.setf(ios::fixed);
  cout.precision(20);
#endif

  int cur = 0;
  for (int i = 0; i < N; ++i) {
    while ((1 << (cur + 1)) < i)
      cur++;
    pw[i] = cur;
  }

  int k;
  cin >> k;
  string s;
  cin >> s;

  int n = s.size();

  k++;

  // split into k parts

  int small = n / k;
  int big = small + 1;
  int cnt_big = n % k;

  if (n % k == 0) {
    vector<string> a;
    for (int i = 0; i + small - 1 < n; i += small) {
      a.app(s.substr(i, small));
    }
    sort(all(a));
    cout << a.back() << endl;
    exit(0);
  }

  vector<int> a;
  for (char c : s)
    a.app(c - '0' + 1);
  a.app(0);

  auto sa = SA.buildSA(a);
  auto lc = SA.buildLCP(a, sa);

  vector<int> pos(n + 1);
  for (int i = 0; i <= n; ++i)
    pos[sa[i]] = i;

  for (int i = 0; i < lc.size(); ++i)
    sp[0][i] = lc[i];
  for (int p = 1; p < LG; ++p) {
    for (int i = 0; i + (1 << p) - 1 < lc.size(); ++i) {
      sp[p][i] = min(sp[p - 1][i], sp[p - 1][i + (1 << (p - 1))]);
    }
  }

  function<int(int, int)> getmin = [&](int l, int r) {
    int len = r - l + 1;
    int p = pw[len];
    return min(sp[p][l], sp[p][r - (1 << p) + 1]);
  };

  auto lcp = [&](int i, int j) {
    i = pos[i];
    j = pos[j];
    if (j < i)
      swap(i, j);

    int ans = N;
    for (int k = i; k < j; ++k)
      ans = min(ans, lc[k]);
    return ans;
  };

  auto ls = [&](int i, int j) {
    int l = lcp(i, j);
    if (l >= big)
      return false;
    else
      return s[i + l] < s[j + l];
  };

  vector<int> ord;
  for (int i = 0; i + big - 1 < n; ++i)
    ord.app(i);
  sort(all(ord), ls);

  function<void(int &, int)> rel = [&](int &a, int b) { a = min(a, b); };

  int l = -1;
  int r = ord.size();
  while (l < r - 1) {
    int m = (l + r) >> 1;

    // all strings <= ord[r]

    for (int i = 1; i < N; ++i)
      dp[i] = N;
    for (int i = 0; i < n; ++i) {
      if (dp[i] == N) {
        continue;
      }

      if (i + small <= n) {
        rel(dp[i + small], dp[i] + 1);
      }
      if (i + big <= n && !ls(ord[m], i)) {
        rel(dp[i + big], dp[i] + 1);
      }
    }
    if (dp[n] <= k)
      r = m;
    else
      l = m;
  }
  cout << s.substr(ord[r], big) << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define ii pair<int, int>
#define app push_back
#define all(a) a.begin(), a.end()
#define bp __builtin_popcountll
#define ll long long
#define mp make_pair
#define f first
#define s second
#define Time (double)clock() / CLOCKS_PER_SEC
#define debug(x) std::cout << #x << ": " << x << '\n';

// It's strictly forbidden to use buildSA() and not
// to add a fictitious integer, which should be smaller
// than all other integers representing a string, to the end!!!
struct SuffixArray {
  static const int SZ = 3e5;

  int c[SZ];
  int cnt[SZ];
  int p[SZ];
  int pn[SZ];
  int cn[SZ];

  vector<int> buildSA(const vector<int> &s) {
    int n = s.size();
    int alpha = (*max_element(s.begin(), s.end())) + 1;
    memset(cnt, 0, alpha * sizeof(int));
    for (int c : s)
      ++cnt[c];
    for (int i = 1; i < alpha; ++i)
      cnt[i] += cnt[i - 1];
    for (int i = 0; i < n; ++i)
      p[--cnt[s[i]]] = i;
    c[p[0]] = 0;
    int cs = 1;
    for (int i = 1; i < n; ++i) {
      if (s[p[i]] != s[p[i - 1]])
        ++cs;
      c[p[i]] = cs - 1;
    }

    for (int h = 0; (1 << h) < n; ++h) {
      for (int i = 0; i < n; ++i) {
        pn[i] = p[i] - (1 << h);
        if (pn[i] < 0)
          pn[i] += n;
      }
      memset(cnt, 0, cs * sizeof(int));
      for (int i = 0; i < n; ++i)
        ++cnt[c[pn[i]]];
      for (int i = 1; i < cs; ++i)
        cnt[i] += cnt[i - 1];
      for (int i = n - 1; i >= 0; --i)
        p[--cnt[c[pn[i]]]] = pn[i];
      cn[p[0]] = 0;
      cs = 1;
      for (int i = 1; i < n; ++i) {
        int mid1 = (p[i] + (1 << h)) % n, mid2 = (p[i - 1] + (1 << h)) % n;
        if (c[p[i]] != c[p[i - 1]] || c[mid1] != c[mid2])
          ++cs;
        cn[p[i]] = cs - 1;
      }
      memcpy(c, cn, n * sizeof(int));
    }

    vector<int> result(p, p + n);
    return result;
  }

  // suf = sa from func above
  vector<int> buildLCP(const vector<int> &s, const vector<int> &suf) const {
    int n = s.size();
    vector<int> rsuf(n);
    vector<int> lcp(n);
    for (ll i = 0; i < n; i++) {
      rsuf[suf[i]] = i;
    }

    int k = 0;
    for (int i = 0; i < n; ++i) {
      if (k > 0)
        --k;
      if (rsuf[i] == n - 1) {
        lcp[n - 1] = -1;
        k = 0;
        continue;
      } else {
        int j = suf[rsuf[i] + 1];
        while (max(i + k, j + k) < n && s[i + k] == s[j + k])
          ++k;
        lcp[rsuf[i]] = k;
      }
    }

    return lcp;
  }
} SA;

const int N = 1e5 + 7, LG = 20;
int dp[N];
int sp[LG][N];

int pw[N];

signed main() {
#ifdef HOME
  freopen("input.txt", "r", stdin);
#else
#define endl '\n'
  ios_base::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  cout.setf(ios::fixed);
  cout.precision(20);
#endif

  int cur = 0;
  for (int i = 0; i < N; ++i) {
    while ((1 << (cur + 1)) < i)
      cur++;
    pw[i] = cur;
  }

  int k;
  cin >> k;
  string s;
  cin >> s;

  int n = s.size();

  k++;

  // split into k parts

  int small = n / k;
  int big = small + 1;
  int cnt_big = n % k;

  if (n % k == 0) {
    vector<string> a;
    for (int i = 0; i + small - 1 < n; i += small) {
      a.app(s.substr(i, small));
    }
    sort(all(a));
    cout << a.back() << endl;
    exit(0);
  }

  vector<int> a;
  for (char c : s)
    a.app(c - '0' + 1);
  a.app(0);

  auto sa = SA.buildSA(a);
  auto lc = SA.buildLCP(a, sa);

  vector<int> pos(n + 1);
  for (int i = 0; i <= n; ++i)
    pos[sa[i]] = i;

  for (int i = 0; i < lc.size(); ++i)
    sp[0][i] = lc[i];
  for (int p = 1; p < LG; ++p) {
    for (int i = 0; i + (1 << p) - 1 < lc.size(); ++i) {
      sp[p][i] = min(sp[p - 1][i], sp[p - 1][i + (1 << (p - 1))]);
    }
  }

  function<int(int, int)> getmin = [&](int l, int r) {
    int len = r - l + 1;
    int p = pw[len];
    return min(sp[p][l], sp[p][r - (1 << p) + 1]);
  };

  auto lcp = [&](int i, int j) {
    i = pos[i];
    j = pos[j];
    if (j < i)
      swap(i, j);

    if (i == j)
      return N;
    else
      return getmin(i, j - 1);
  };

  auto ls = [&](int i, int j) {
    int l = lcp(i, j);
    if (l >= big)
      return false;
    else
      return s[i + l] < s[j + l];
  };

  vector<int> ord;
  for (int i = 0; i + big - 1 < n; ++i)
    ord.app(i);
  sort(all(ord), ls);

  function<void(int &, int)> rel = [&](int &a, int b) { a = min(a, b); };

  int l = -1;
  int r = ord.size();
  while (l < r - 1) {
    int m = (l + r) >> 1;

    // all strings <= ord[r]

    for (int i = 1; i < N; ++i)
      dp[i] = N;
    for (int i = 0; i < n; ++i) {
      if (dp[i] == N) {
        continue;
      }

      if (i + small <= n) {
        rel(dp[i + small], dp[i] + 1);
      }
      if (i + big <= n && !ls(ord[m], i)) {
        rel(dp[i + big], dp[i] + 1);
      }
    }
    if (dp[n] <= k)
      r = m;
    else
      l = m;
  }
  cout << s.substr(ord[r], big) << endl;
}