2022-12-15

DISCO presents 2016 予選 C - アメージングな文字列は、きみが作る！ (橙色)

AtCoder unrated公式コン橙色diff SuffixArray 場合分け辞書順 Greedy prefixとsuffix 文字列操作操作:挿入操作:1文字を変更する操作後の結果の最適化問題最適化テク:操作の流れを単純化する操作をちょうどK回行う思わず解きたくなる興味深い良問そのまま覚えたいシンプル設定の中堅以上の典型問題 Greedy:辞書順最小を求める

とても面白かった。文字列に操作を $K$ 回施して、操作後の文字列の辞書順最小のものを求める問題。Suffix Array のよい練習問題でもある。

問題へのリンク

問題概要

英小文字のみからなる長さ $N$ の文字列 $S$ が与えられる。この文字列に対して、以下のいずれかの作業をちょうど $K$ 回実行する

文字列 $S$ の好きな箇所に好きな英小文字を挿入する
文字列 $S$ 中の好きな文字を削除する
文字列 $S$ の好きな文字を好きな英小文字に変更する

できあがる文字列のうち、辞書順最小な文字列を答えよ。

制約

$2 \le N \le 3 \times 10^{5}$
$1 \le K \le N - 1$

考えたこと

まずは様子を掴んでみる

$S$ = "abc", $K = 1$ → "aabc"
$S$ = "acb", $K = 1$ → "aab"
$S$ = "aba", $K = 1$ → "aa"
...

やっていくうちに、'a' だけにしてしまった方がいいやつと、'a' を先頭に多くした方がいいやつがあるらしいことに気がついた。

'a' だけにできるやつ

文字列 $S$ 中の 'a' 以外の文字の個数を $m$ としたとき、 $m \le K$ ならば 'a' だけにできる。'a' だけにできるならば、文字数は少ない方がいいので $N-K$ 文字にする。つまり、答えは string(N-K, 'a') になる。

なお、 $K \le N -1$ なので空文字列にはできないことに注意する。

'a' だけにできないやつ

これはどうしようもないので、先頭に 'a' を付け加えて引き伸ばすことにする。しかし、すでにある文字を 'a' に変更した方が良い場合もある。そこで、次のように考えることにした。

先頭から $i$ 文字分については、'a' でない文字をすべて 'a' に変更した上で、あまった操作回数は 'a' を先頭に挿入していく
$i$ 文字目以降は何もしない

この場合だけ考えればよいことは比較的すぐにわかった。'a' 以外の文字をそのままにした場合、それ以降の文字を削除したり変更したりする意味はないのだ。

あとは、 $i = 0, 1, \dots$ に対して、辞書順最小のものを求めればよい (ただし 'a' 以外の文字が $K$ 以内の範囲のみ)。

比較方法は次のように考えることにした。

'a' の個数が最大のもの、ただし 'a' の個数が等しい場合は、そのうちの S.substr(i) の辞書順が最小のもの

S.substr(i) の辞書順については Suffix Array を用いることで比較できる。

全体として、計算量は $O(N)$ となる。

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

// SA-IS (O(N))
template<class Str> struct SuffixArray {
    // data
    Str str;
    vector<int> sa;   // sa[i] : the starting index of the i-th smallest suffix (i = 0, 1, ..., n)
    vector<int> lcp;  // lcp[i]: the lcp of sa[i] and sa[i+1] (i = 0, 1, ..., n-1)
    int& operator [] (int i) {
        return sa[i];
    }

    // constructor
    SuffixArray(const Str& str_) : str(str_) {
        build_sa();
    }
    void init(const Str& str_) {
        str = str_;
        build_sa();
    }
    void build_sa() {
        int N = (int)str.size();
        vector<int> s;
        for (int i = 0; i < N; ++i) s.push_back(str[i] + 1);
        s.push_back(0);
        sa = sa_is(s);
        calcLCP(s);
    }

    // SA-IS
    // upper: # of characters 
    vector<int> sa_is(vector<int> &s, int upper = 256) {
        int N = (int)s.size();
        if (N == 0) return {};
        else if (N == 1) return {0};
        else if (N == 2) {
            if (s[0] < s[1]) return {0, 1};
            else return {1, 0};
        }

        vector<int> isa(N);
        vector<bool> ls(N, false);
        for (int i = N - 2; i >= 0; --i) {
            ls[i] = (s[i] == s[i + 1]) ? ls[i + 1] : (s[i] < s[i + 1]);
        }
        vector<int> sum_l(upper + 1, 0), sum_s(upper + 1, 0);
        for (int i = 0; i < N; ++i) {
            if (!ls[i]) ++sum_s[s[i]];
            else ++sum_l[s[i] + 1];
        }
        for (int i = 0; i <= upper; ++i) {
            sum_s[i] += sum_l[i];
            if (i < upper) sum_l[i + 1] += sum_s[i];
        }

        auto induce = [&](const vector<int> &lms) -> void {
            fill(isa.begin(), isa.end(), -1);
            vector<int> buf(upper + 1);
            copy(sum_s.begin(), sum_s.end(), buf.begin());
            for (auto d: lms) {
                if (d == N) continue;
                isa[buf[s[d]]++] = d;
            }
            copy(sum_l.begin(), sum_l.end(), buf.begin());
            isa[buf[s[N - 1]]++] = N - 1;
            for (int i = 0; i < N; ++i) {
                int v = isa[i];
                if (v >= 1 && !ls[v - 1]) {
                    isa[buf[s[v - 1]]++] = v - 1;
                }
            }
            copy(sum_l.begin(), sum_l.end(), buf.begin());
            for (int i = N - 1; i >= 0; --i) {
                int v = isa[i];
                if (v >= 1 && ls[v - 1]) {
                    isa[--buf[s[v - 1] + 1]] = v - 1;
                }
            }
        };
            
        vector<int> lms, lms_map(N + 1, -1);
        int M = 0;
        for (int i = 1; i < N; ++i) {
            if (!ls[i - 1] && ls[i]) {
                lms_map[i] = M++;
            }
        }
        lms.reserve(M);
        for (int i = 1; i < N; ++i) {
            if (!ls[i - 1] && ls[i]) {
                lms.push_back(i);
            }
        }
        induce(lms);

        if (M) {
            vector<int> lms2;
            lms2.reserve(isa.size());
            for (auto v: isa) {
                if (lms_map[v] != -1) lms2.push_back(v);
            }
            int rec_upper = 0;
            vector<int> rec_s(M);
            rec_s[lms_map[lms2[0]]] = 0;
            for (int i = 1; i < M; ++i) {
                int l = lms2[i - 1], r = lms2[i];
                int nl = (lms_map[l] + 1 < M) ? lms[lms_map[l] + 1] : N;
                int nr = (lms_map[r] + 1 < M) ? lms[lms_map[r] + 1] : N;
                bool same = true;
                if (nl - l != nr - r) same = false;
                else {
                    while (l < nl) {
                        if (s[l] != s[r]) break;
                        ++l, ++r;
                    }
                    if (l == N || s[l] != s[r]) same = false;
                }
                if (!same) ++rec_upper;
                rec_s[lms_map[lms2[i]]] = rec_upper;
            }
            auto rec_sa = sa_is(rec_s, rec_upper);

            vector<int> sorted_lms(M);
            for (int i = 0; i < M; ++i) {
                sorted_lms[i] = lms[rec_sa[i]];
            }
            induce(sorted_lms);
        }
        return isa;
    }

    // prepair lcp
    vector<int> rsa;
    void calcLCP(const vector<int> &s) {
        int N = (int)s.size();
        rsa.assign(N, 0), lcp.assign(N, 0);
        for (int i = 0; i < N; ++i) rsa[sa[i]] = i;
        int h = 0;
        for (int i = 0; i < N - 1; ++i) {
            int pi = sa[rsa[i] - 1];
            if (h > 0) --h;
            for (; pi + h < N && i + h < N; ++h) {
                if (s[pi + h] != s[i + h]) break;
            }
            lcp[rsa[i] - 1] = h;
        }
    }
};

int main() {
    int K;
    string S;
    cin >> S >> K;
    int N = S.size();

    // Suffix Array を構築して rsa を取得
    SuffixArray<string> sa(S);
    vector<int> rsa = sa.rsa;

    // a だけにできるか
    int nota = 0;
    for (auto c : S) if (c != 'a') ++nota;
    if (nota <= K) {
        cout << string(N - K, 'a') << endl;
        return 0;
    }

    nota = 0;
    pair<int,int> mi = {N, N};
    int resid = -1;
    for (int i = 0; i < N; ++i) {
        if (nota > K) break;
        int num = i + (K - nota);
        int ord = rsa[i];
        if (chmin(mi, make_pair(-num, ord))) resid = i;
        if (S[i] != 'a') ++nota;
    }
    cout << string(-mi.first, 'a') + S.substr(resid) << endl;
}