Solution

C.cpp

#include <iostream>
#include <cstdint>
#include <vector>
#include <string>
#include <fstream>
#include <set>

using namespace std;

int main() {
    ifstream in("shifts.in");
    ofstream out("shifts.out");

    string s;
    getline(in, s);

    int64_t id;
    in >> id;

    int64_t k = '~' - ' ' + 1;
    int64_t n = s.size();

    // Calc color
    vector<int64_t> cnt(max(n, k)), col(n);

    for(int64_t i = 0; i < n; i++) {
        col[i] = s[i] - ' ';
        cnt[col[i]]++;
    }

    // Have references to beggining of each block
    vector<int64_t> head(max(n, k));

    head[0] = 0;
    for(int64_t i = 1; i < k; i++) {
        head[i] = head[i - 1] + cnt[i - 1];
    }

    // Count sort
    vector<int64_t> a(n);
    for(int64_t i = 0; i < n; i++) {
        a[head[col[i]]++] = i;
    }

    // Fix references to beggining of each block
    head[0] = 0;
    for(int64_t i = 1; i < k; i++) {
        head[i] = head[i - 1] + cnt[i - 1];
    }

    for(int64_t l = 1; l < n; l *= 2) {
        // e.g. a = |10|11|2 |4 |6 |9 |0 |7 |3 |6|1 |8 |

        // Resort based on the fact that B is sorted
        vector<int64_t> a_(n);

        for(int64_t i = 0; i < n; i++) {
            int64_t j = (a[i] + 2 * n - l) % n;
            a_[head[col[j]]++] = j;
        }

        // Recalc colors
        vector<int64_t> col_(n);

        k = 0;

        int64_t k = 1;
        for(int64_t i = 0; i < n; i++) {
            a[i] = a_[i];
            if(i == 0 || col[a[i]] != col[a[i - 1]] || col[(a[i] + l) % n] != col[(a[i - 1] + l) % n]) {
                if(k >= head.size()) {
                    head.resize(k + 1);
                }
                head[k++] = i;
            }
            col_[a[i]] = k - 1;
        }

        col = col_;
    }

    // Kosai

    vector<int64_t> ord(n), lcp(n);
    for(int64_t i = 0; i < n; i++) {
        ord[a[i]] = i;
    }

    int64_t x = 0;
    for(int64_t i = 0; i < n; i++) {
        x = max(x - 1, 0LL);
        if(ord[i] == n - 1) {
            x = 0;
        } else {
            int64_t j = a[ord[i] + 1];
            while(s[i + x] == s[j + x]) {
                x++;
            }
            lcp[ord[i]] = x;
        }
    }

    int64_t last_color = 0;
    int64_t color_cnt = 0;

    for(int64_t i = 0; i < col.size(); i++) {
        if(col[a[i]] != last_color) {
            last_color = col[a[i]];
            color_cnt++;
            if(color_cnt == id) {
                out << s.substr(a[i]) << s.substr(0, a[i]) << endl;
                return 0;
            }
        }
    }
    out << "IMPOSSIBLE";

    return 0;
}