Problem p03355 (C++)

Buggy submission (status: Time Limit Exceeded):
#include <algorithm>
#include <cmath>
#include <complex> // complex<double> a(1.2 , 2.3);// real(): 1.2, imag()2.3
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <vector>
using namespace std;

#define MOD 1000000007
#define ll long long
#define ld long double
#define FOR(i, a, b) for (ll i = (ll)a; i < (ll)b; i++)
#define rep(i, n) FOR(i, 0, n)
#define pb push_back
#define mp make_pair
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define rmsame(a) sort(all(a)), a.erase(unique(all(a)), a.end())
#define rmvector(a, b)                                                         \
  rep(i, a.size()) rep(j, b.size()) if (a[i] == b[j]) {                        \
    a.erase(a.begin() + i);                                                    \
    i--;                                                                       \
    break;                                                                     \
  }
#define pq_pair_tB                                                             \
  priority_queue<pair<ll, ll>, vector<pair<ll, ll>>, greater<pair<ll, ll>>>
#define pq_pair_ts priority_queue<pair<ll, ll>> // 第二成分の大きさが関係ない
template <typename X> bool exist(vector<X> vec, X item) {
  return find(all(vec), item) != vec.end();
}
ll gcd(ll a, ll b) {
  if (b == 0)
    return a;
  return gcd(b, a % b);
}
ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }

vector<string> ans[26];

ll K;
void cal(ll index) {
  if (K == 0)
    return;
  vector<string> ret;
  rep(i, ans[index].size()) {
    rmsame(ans[index]);
    rep(j, (ans[index][i].size())) { ret.pb(ans[index][i].substr(0, j + 1)); }
  }
  rmsame(ret);
  if (K > ret.size())
    K -= ret.size();
  else {
    cout << ret[K - 1] << endl;
    K = 0;
  }
}

string find(ll index, ll num) {
  vector<string> ret;
  rep(i, ans[index].size()) {
    rmsame(ans[index]);
    rep(j, (ans[index][i].size())) { ret.pb(ans[index][i].substr(0, j + 1)); }
  }
  rmsame(ret);
  return ret[num - 1];
}
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  string s;
  cin >> s;
  cin >> K;

  rep(i, s.size()) {
    // cout << s.substr(i,5) << endl;
    ans[s[i] - 'a'].pb(s.substr(i));
  }
  //  cout <<"--" << endl;
  rep(i, 26) { cal(i); }

  // cout << fixed << setprecision(16) << ans << endl;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cmath>
#include <complex> // complex<double> a(1.2 , 2.3);// real(): 1.2, imag()2.3
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <vector>
using namespace std;

#define MOD 1000000007
#define ll long long
#define ld long double
#define FOR(i, a, b) for (ll i = (ll)a; i < (ll)b; i++)
#define rep(i, n) FOR(i, 0, n)
#define pb push_back
#define mp make_pair
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define rmsame(a) sort(all(a)), a.erase(unique(all(a)), a.end())
#define rmvector(a, b)                                                         \
  rep(i, a.size()) rep(j, b.size()) if (a[i] == b[j]) {                        \
    a.erase(a.begin() + i);                                                    \
    i--;                                                                       \
    break;                                                                     \
  }
#define pq_pair_tB                                                             \
  priority_queue<pair<ll, ll>, vector<pair<ll, ll>>, greater<pair<ll, ll>>>
#define pq_pair_ts priority_queue<pair<ll, ll>> // 第二成分の大きさが関係ない
template <typename X> bool exist(vector<X> vec, X item) {
  return find(all(vec), item) != vec.end();
}
ll gcd(ll a, ll b) {
  if (b == 0)
    return a;
  return gcd(b, a % b);
}
ll lcm(ll a, ll b) { return a / gcd(a, b) * b; }

vector<string> ans[26];

ll K;
void cal(ll index) {
  if (K == 0)
    return;
  vector<string> ret;
  rep(i, ans[index].size()) {
    rmsame(ans[index]);
    rep(j, (ans[index][i].size())) { ret.pb(ans[index][i].substr(0, j + 1)); }
  }
  rmsame(ret);
  if (K > ret.size())
    K -= ret.size();
  else {
    cout << ret[K - 1] << endl;
    K = 0;
  }
}

string find(ll index, ll num) {
  vector<string> ret;
  rep(i, ans[index].size()) {
    rmsame(ans[index]);
    rep(j, (ans[index][i].size())) { ret.pb(ans[index][i].substr(0, j + 1)); }
  }
  rmsame(ret);
  return ret[num - 1];
}
int main() {
  cin.tie(0);
  ios::sync_with_stdio(false);

  string s;
  cin >> s;
  cin >> K;

  rep(i, s.size()) {
    // cout << s.substr(i,5) << endl;
    ans[s[i] - 'a'].pb(s.substr(i, 5));
  }
  //  cout <<"--" << endl;
  rep(i, 26) { cal(i); }

  // cout << fixed << setprecision(16) << ans << endl;
  return 0;
}