Problem p02131 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>

using namespace std;

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

  void Build_SA(const string &str) {
    s = str;
    SA.resize(s.size());
    iota(begin(SA), end(SA), 0);
    sort(begin(SA), end(SA), [&](const int &a, const int &b) {
      if (s[a] == s[b])
        return (a > b);
      return (s[a] < s[b]);
    });
    vector<int> classes(s.size()), c(s.size()), cnt(s.size());
    for (int i = 0; i < s.size(); i++) {
      c[i] = s[i];
    }
    for (int len = 1; len < s.size(); len <<= 1) {
      for (int i = 0; i < s.size(); i++) {
        if (i > 0 && c[SA[i - 1]] == c[SA[i]] && SA[i - 1] + len < s.size() &&
            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 < 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 (s.size()); }

  bool lt_substr(string &t, int si = 0, int ti = 0) {
    int sn = s.size(), tn = t.size();
    while (si < sn && ti < tn) {
      if (s[si] < t[ti])
        return (true);
      if (s[si] > t[ti])
        return (false);
      ++si, ++ti;
    }
    return (si >= sn && ti < tn);
  }

  int lower_bound(string &t) {
    int low = -1, high = SA.size();
    while (high - low > 1) {
      int mid = (low + high) >> 1;
      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 = SA.size();
    t.back()++;
    while (high - low > 1) {
      int mid = (low + high) >> 1;
      if (lt_substr(t, SA[mid]))
        low = mid;
      else
        high = mid;
    }
    t.back()--;
    return (make_pair(idx, high));
  }

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

using int64 = long long;

template <class T> struct BinaryIndexedTree {
  vector<T> data;

  BinaryIndexedTree(int sz) { data.assign(++sz, 0); }

  T sum(int k) {
    T ret = 0;
    for (++k; k > 0; k -= k & -k)
      ret += data[k];
    return (ret);
  }

  void add(int k, T x) {
    for (++k; k < data.size(); k += k & -k)
      data[k] += x;
  }
};

template <class T> struct AdditionalBIT : BinaryIndexedTree<T> {
  int curr;

  AdditionalBIT(int sz) : BinaryIndexedTree<T>(sz) {
    curr = 1;
    while (curr <= sz)
      curr <<= 1;
  }

  int lower_bound(T w) {
    if (w <= 0)
      return (0);
    int i = 0;
    for (int k = curr; k > 0; k >>= 1) {
      if (i + k < BinaryIndexedTree<T>::data.size() &&
          BinaryIndexedTree<T>::data[i + k] < w) {
        w -= BinaryIndexedTree<T>::data[i + k];
        i += k;
      }
    }
    return (i);
  }
};

int main() {
  int N, H;
  string S;

  cin >> N >> H;
  cin >> S;

  int wait = 0;
  vector<int> weight(S.size());
  for (int i = 0; i < S.size(); i++) {
    if (S[i] == '0') {
      ++wait;
      H -= wait;
    } else {
      weight[i] += wait + 1;
      wait = 0;
    }
  }

  SuffixArray sa;
  sa.Build_SA(S);

  if (H <= 0) {
    cout << 0 << endl;
    return (0);
  }

  vector<int> order;
  for (int i = 0; i < S.size(); i++) {
    if (S[sa[i]] == '0')
      continue;
    order.emplace_back(i);
  }

  AdditionalBIT<int64> bit(N);
  vector<int> end[100001];
  for (auto &v : order) {
    bit.add(v, weight[sa[v]]);
    end[N - sa[v]].emplace_back(v);
  }

  int keta = 1;
  for (; H - bit.sum(N - 1) >= 1; ++keta) {
    H -= bit.sum(N - 1);
    for (auto &p : end[keta])
      bit.add(p, -weight[sa[p]]);
  }

  cout << S.substr(sa[bit.lower_bound(H)], keta) << endl;
}


Fixed submission:
#include <bits/stdc++.h>

using namespace std;

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

  void Build_SA(const string &str) {
    s = str;
    SA.resize(s.size());
    iota(begin(SA), end(SA), 0);
    sort(begin(SA), end(SA), [&](const int &a, const int &b) {
      if (s[a] == s[b])
        return (a > b);
      return (s[a] < s[b]);
    });
    vector<int> classes(s.size()), c(s.size()), cnt(s.size());
    for (int i = 0; i < s.size(); i++) {
      c[i] = s[i];
    }
    for (int len = 1; len < s.size(); len <<= 1) {
      for (int i = 0; i < s.size(); i++) {
        if (i > 0 && c[SA[i - 1]] == c[SA[i]] && SA[i - 1] + len < s.size() &&
            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 < 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 (s.size()); }

  bool lt_substr(string &t, int si = 0, int ti = 0) {
    int sn = s.size(), tn = t.size();
    while (si < sn && ti < tn) {
      if (s[si] < t[ti])
        return (true);
      if (s[si] > t[ti])
        return (false);
      ++si, ++ti;
    }
    return (si >= sn && ti < tn);
  }

  int lower_bound(string &t) {
    int low = -1, high = SA.size();
    while (high - low > 1) {
      int mid = (low + high) >> 1;
      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 = SA.size();
    t.back()++;
    while (high - low > 1) {
      int mid = (low + high) >> 1;
      if (lt_substr(t, SA[mid]))
        low = mid;
      else
        high = mid;
    }
    t.back()--;
    return (make_pair(idx, high));
  }

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

using int64 = long long;

template <class T> struct BinaryIndexedTree {
  vector<T> data;

  BinaryIndexedTree(int sz) { data.assign(++sz, 0); }

  T sum(int k) {
    T ret = 0;
    for (++k; k > 0; k -= k & -k)
      ret += data[k];
    return (ret);
  }

  void add(int k, T x) {
    for (++k; k < data.size(); k += k & -k)
      data[k] += x;
  }
};

template <class T> struct AdditionalBIT : BinaryIndexedTree<T> {
  int curr;

  AdditionalBIT(int sz) : BinaryIndexedTree<T>(sz) {
    curr = 1;
    while (curr <= sz)
      curr <<= 1;
  }

  int lower_bound(T w) {
    if (w <= 0)
      return (0);
    int i = 0;
    for (int k = curr; k > 0; k >>= 1) {
      if (i + k < BinaryIndexedTree<T>::data.size() &&
          BinaryIndexedTree<T>::data[i + k] < w) {
        w -= BinaryIndexedTree<T>::data[i + k];
        i += k;
      }
    }
    return (i);
  }
};

int main() {
  int64 N, H;
  string S;

  cin >> N >> H;
  cin >> S;

  int wait = 0;
  vector<int> weight(S.size());
  for (int i = 0; i < S.size(); i++) {
    if (S[i] == '0') {
      ++wait;
      H -= wait;
    } else {
      weight[i] += wait + 1;
      wait = 0;
    }
  }

  SuffixArray sa;
  sa.Build_SA(S);

  if (H <= 0) {
    cout << 0 << endl;
    return (0);
  }

  vector<int> order;
  for (int i = 0; i < S.size(); i++) {
    if (S[sa[i]] == '0')
      continue;
    order.emplace_back(i);
  }

  AdditionalBIT<int64> bit(N);
  vector<int> end[100001];
  for (auto &v : order) {
    bit.add(v, weight[sa[v]]);
    end[N - sa[v]].emplace_back(v);
  }

  int keta = 1;
  for (; H - bit.sum(N - 1) >= 1; ++keta) {
    H -= bit.sum(N - 1);
    for (auto &p : end[keta])
      bit.add(p, -weight[sa[p]]);
  }

  cout << S.substr(sa[bit.lower_bound(H)], keta) << endl;
}