Problem p03904 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <bits/stdc++.h>
using namespace std;
#define ALL(x) x.begin(), x.end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define debug(v)                                                               \
  cout << #v << ":";                                                           \
  for (auto x : v) {                                                           \
    cout << x << ' ';                                                          \
  }                                                                            \
  cout << endl;
#define mod 1000000007
using ll = long long;
const int INF = 1000000000;
const ll LINF = 1001002003004005006ll;
int dx[] = {1, 0, -1, 0}, dy[] = {0, 1, 0, -1};
// ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

struct IOSetup {
  IOSetup() {
    cin.tie(0);
    ios::sync_with_stdio(0);
    cout << fixed << setprecision(12);
  }
} iosetup;

template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  os << p.first << " " << p.second;
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++)
    os << v[i] << (i + 1 == (int)v.size() ? "" : " ");
  return os;
}

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &x : v)
    is >> x;
  return is;
}

struct SuffixArray {
  vector<int> SA;
  const string s;

  SuffixArray(const string &str) : s(str) {
    SA.resize(s.size());
    iota(begin(SA), end(SA), 0);
    sort(begin(SA), end(SA),
         [&](int a, int b) { return s[a] == s[b] ? a > b : s[a] < s[b]; });
    vector<int> classes(s.size()), c(s.begin(), s.end()), cnt(s.size());
    for (int len = 1; len < (int)s.size(); len <<= 1) {
      for (int i = 0; i < (int)s.size(); i++) {
        if (i > 0 and c[SA[i - 1]] == c[SA[i]] and
            SA[i - 1] + len < (int)s.size() and
            c[SA[i - 1] + len / 2] == c[SA[i] + len / 2]) {
          classes[SA[i]] = classes[SA[i - 1]];
        } else {
          classes[SA[i]] = i;
        }
      }
      iota(begin(cnt), end(cnt), 0);
      copy(begin(SA), end(SA), begin(c));
      for (int i = 0; i < (int)s.size(); i++) {
        int s1 = c[i] - len;
        if (s1 >= 0)
          SA[cnt[classes[s1]]++] = s1;
      }
      classes.swap(c);
    }
  }
  int operator[](int k) const { return SA[k]; }
  int size() const { return (int)s.size(); }
  bool lt_substr(const string &t, int si = 0, int ti = 0) {
    int sn = (int)s.size(), tn = (int)t.size();
    while (si < sn and ti < tn) {
      if (s[si] < t[ti])
        return true;
      if (s[si] > t[ti])
        return false;
      ++si, ++ti;
    }
    return si >= sn and ti < tn;
  }

  int lower_bound(const string &t) {
    int low = -1, high = (int)SA.size();
    while (high - low > 1) {
      int mid = (low + high) / 2;
      if (lt_substr(t, SA[mid]))
        low = mid;
      else
        high = mid;
    }
    return high;
  }

  pair<int, int> lower_upper_bound(string &t) {
    int idx = lower_bound(t);
    int low = idx - 1, high = (int)SA.size();
    t.back()++;
    while (high - low > 1) {
      int mid = (low + high) / 2;
      if (lt_substr(t, SA[mid]))
        low = mid;
      else
        high = mid;
    }
    t.back()--;
    return {idx, high};
  }

  void output() {
    for (int i = 0; i < size(); i++) {
      cout << i << " : " << s.substr(SA[i]) << endl;
    }
  }
};

struct LongestCommonPrefixArray {
  const SuffixArray &SA;
  vector<int> LCP, rank;

  LongestCommonPrefixArray(const SuffixArray &SA) : SA(SA), LCP(SA.size()) {
    rank.resize(SA.size());
    for (int i = 0; i < SA.size(); i++) {
      rank[SA[i]] = i;
    }
    for (int i = 0, h = 0; i < (int)SA.size(); i++) {
      if (rank[i] + 1 < (int)SA.size()) {
        for (int j = SA[rank[i] + 1];
             max(i, j) + h < (int)SA.size() and SA.s[i + h] == SA.s[j + h]; ++h)
          ;
        LCP[rank[i] + 1] = h;
        if (h > 0)
          --h;
      }
    }
  }

  int operator[](int k) const { return LCP[k]; }

  size_t size() const { return LCP.size(); }

  void output() {
    for (int i = 0; i < (int)size(); i++) {
      cout << i << ": " << LCP[i] << " " << SA.s.substr(SA[i]) << endl;
    }
  }
};

signed main() {
  int k;
  cin >> k;
  k++;
  string s;
  cin >> s;
  int n = (int)s.size();

  SuffixArray sa(s);
  vector<int> rank(sa.size());
  rep(i, (int)sa.size()) rank[sa[i]] = i;

  // t桁にする．
  int t = (n + k - 1) / k;
  // p個はt-1桁でいいよ
  // int p=t*k-n;

  // 候補
  vector<int> cand;
  rep(i, n) {
    int len = n - sa[i];
    if (len >= t)
      cand.push_back(i);
  }

  // sa.output();
  // debug(cand);
  // debug(rank);

  // sa[m]~sa[m]+tのやつが答えになるか？
  auto check = [&](int m) {
    int idx;
    int targetrank = cand[m];
    int cnt = 0;
    for (idx = 0; idx < n;) {
      cnt++;
      if (n - idx < t)
        break;
      if (rank[idx] <= targetrank)
        idx += t;
      else
        idx += t - 1;
    }
    return cnt <= k;
  };

  int lw = 0, hi = (int)cand.size() - 1, res = 0;
  while (lw <= hi) {
    int mid = (lw + hi) / 2;
    if (check(mid)) {
      res = mid;
      hi = mid - 1;
    } else {
      lw = mid + 1;
    }
  }
  int i = sa[cand[res]];
  cout << s.substr(i, t) << endl;
  return 0;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define ALL(x) x.begin(), x.end()
#define rep(i, n) for (int i = 0; i < (n); i++)
#define debug(v)                                                               \
  cout << #v << ":";                                                           \
  for (auto x : v) {                                                           \
    cout << x << ' ';                                                          \
  }                                                                            \
  cout << endl;
#define mod 1000000007
using ll = long long;
const int INF = 1000000000;
const ll LINF = 1001002003004005006ll;
int dx[] = {1, 0, -1, 0}, dy[] = {0, 1, 0, -1};
// ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
template <class T> bool chmax(T &a, const T &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

struct IOSetup {
  IOSetup() {
    cin.tie(0);
    ios::sync_with_stdio(0);
    cout << fixed << setprecision(12);
  }
} iosetup;

template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &p) {
  os << p.first << " " << p.second;
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const vector<T> &v) {
  for (int i = 0; i < (int)v.size(); i++)
    os << v[i] << (i + 1 == (int)v.size() ? "" : " ");
  return os;
}

template <typename T1, typename T2>
istream &operator>>(istream &is, pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T> istream &operator>>(istream &is, vector<T> &v) {
  for (T &x : v)
    is >> x;
  return is;
}

struct SuffixArray {
  vector<int> SA;
  const string s;

  SuffixArray(const string &str) : s(str) {
    SA.resize(s.size());
    iota(begin(SA), end(SA), 0);
    sort(begin(SA), end(SA),
         [&](int a, int b) { return s[a] == s[b] ? a > b : s[a] < s[b]; });
    vector<int> classes(s.size()), c(s.begin(), s.end()), cnt(s.size());
    for (int len = 1; len < (int)s.size(); len <<= 1) {
      for (int i = 0; i < (int)s.size(); i++) {
        if (i > 0 and c[SA[i - 1]] == c[SA[i]] and
            SA[i - 1] + len < (int)s.size() and
            c[SA[i - 1] + len / 2] == c[SA[i] + len / 2]) {
          classes[SA[i]] = classes[SA[i - 1]];
        } else {
          classes[SA[i]] = i;
        }
      }
      iota(begin(cnt), end(cnt), 0);
      copy(begin(SA), end(SA), begin(c));
      for (int i = 0; i < (int)s.size(); i++) {
        int s1 = c[i] - len;
        if (s1 >= 0)
          SA[cnt[classes[s1]]++] = s1;
      }
      classes.swap(c);
    }
  }
  int operator[](int k) const { return SA[k]; }
  int size() const { return (int)s.size(); }
  bool lt_substr(const string &t, int si = 0, int ti = 0) {
    int sn = (int)s.size(), tn = (int)t.size();
    while (si < sn and ti < tn) {
      if (s[si] < t[ti])
        return true;
      if (s[si] > t[ti])
        return false;
      ++si, ++ti;
    }
    return si >= sn and ti < tn;
  }

  int lower_bound(const string &t) {
    int low = -1, high = (int)SA.size();
    while (high - low > 1) {
      int mid = (low + high) / 2;
      if (lt_substr(t, SA[mid]))
        low = mid;
      else
        high = mid;
    }
    return high;
  }

  pair<int, int> lower_upper_bound(string &t) {
    int idx = lower_bound(t);
    int low = idx - 1, high = (int)SA.size();
    t.back()++;
    while (high - low > 1) {
      int mid = (low + high) / 2;
      if (lt_substr(t, SA[mid]))
        low = mid;
      else
        high = mid;
    }
    t.back()--;
    return {idx, high};
  }

  void output() {
    for (int i = 0; i < size(); i++) {
      cout << i << " : " << s.substr(SA[i]) << endl;
    }
  }
};

struct LongestCommonPrefixArray {
  const SuffixArray &SA;
  vector<int> LCP, rank;

  LongestCommonPrefixArray(const SuffixArray &SA) : SA(SA), LCP(SA.size()) {
    rank.resize(SA.size());
    for (int i = 0; i < SA.size(); i++) {
      rank[SA[i]] = i;
    }
    for (int i = 0, h = 0; i < (int)SA.size(); i++) {
      if (rank[i] + 1 < (int)SA.size()) {
        for (int j = SA[rank[i] + 1];
             max(i, j) + h < (int)SA.size() and SA.s[i + h] == SA.s[j + h]; ++h)
          ;
        LCP[rank[i] + 1] = h;
        if (h > 0)
          --h;
      }
    }
  }

  int operator[](int k) const { return LCP[k]; }

  size_t size() const { return LCP.size(); }

  void output() {
    for (int i = 0; i < (int)size(); i++) {
      cout << i << ": " << LCP[i] << " " << SA.s.substr(SA[i]) << endl;
    }
  }
};

signed main() {
  int k;
  cin >> k;
  k++;
  string s;
  cin >> s;
  int n = (int)s.size();

  SuffixArray sa(s);
  vector<int> rank(sa.size());
  rep(i, (int)sa.size()) rank[sa[i]] = i;

  // t桁にする．
  int t = (n + k - 1) / k;
  // p個はt-1桁でいいよ
  // int p=t*k-n;

  // 候補
  vector<int> cand;
  rep(i, n) {
    int len = n - sa[i];
    if (len >= t)
      cand.push_back(i);
  }

  // sa.output();
  // debug(cand);
  // debug(rank);

  // sa[m]~sa[m]+tのやつが答えになるか？
  auto check = [&](int m) {
    int idx;
    int targetrank = cand[m];
    int cnt = 0;
    for (idx = 0; idx < n;) {
      cnt++;
      if (n - idx < t)
        break;
      if (rank[idx] <= targetrank)
        idx += t;
      else
        idx += t - 1;
      if (cnt > k)
        return false;
    }
    return cnt <= k;
  };

  int lw = 0, hi = (int)cand.size() - 1, res = 0;
  while (lw <= hi) {
    int mid = (lw + hi) / 2;
    if (check(mid)) {
      res = mid;
      hi = mid - 1;
    } else {
      lw = mid + 1;
    }
  }
  int i = sa[cand[res]];
  cout << s.substr(i, t) << endl;
  return 0;
}