Codeforces 1493C - K-beautiful Strings (設計解題、實作題)

#include <iostream>
#include <bits/stdc++.h>
#define LOCAL
#define MAXN 100020
using namespace std;
int T, k, n;
int cnt[30], sum[MAXN][30], now;
//cnt 如果在這 index 後的字串後綴都替換後，需要給予的字元數量，index 是 a,b,c 這種感覺
//sum[字串當前長度][字母 (a,b, c)] = 數量，也就是字串當前長度的某個字元的數量
string s; //輸入題目資料用
 
int check(int x){ //判斷在此 index 下能否表達出題目所要的答案字元
    for(int i = s[x]-'a'+1; i < 26; i++){ //從修改的 index 字元的下一個開始，到 z 結束
        memset(cnt, 0, sizeof(cnt)); //將 cnt 歸為零，cnt 每次都要 clear
        //cnt 用來表示，我們的 cnt[字母 (a,b,c)] = 數量，也就是每個字母需要幾個數量才能符合，
        //同個字母數量都能 mod k == 0

        int remain = n-x-1; //如果從此 index 開始修改，那後面的字元我們可以任意變動
        //但也只有這個數量可以變動，注意要 -1，因為 x 這個 index 的字元是固定的，
        //必須要是 i，才有符合當前字源比 s[i] 小

        for(int j = 0; j < 26; j++){ //每個字元還需要多少字元才能符合，
        //同個字母數量都能 mod k == 0
            now = sum[x][j]; //當前長度的 j + 'a' 字元有多少數量
            if(j == s[x]-'a') now--; //由於我們是紀錄當前長度，但是我們要替換現在此 index，
            //因此要將 now-1，因為 s[i] 這個字元我們不會使用，如果使用就不是修改此 index
            else if(i == j) now++; //因為在這邊我們會用 i+'a' 此字元，所以將 now +1，
            //表示之後可以少一個 i+'a' 字元，來補我們這邊用過的

            remain -= (k-now%k)%k; //減掉我們可以用的長度字元
            cnt[j] += (k-now%k)%k; //此字元至少需要多少長度才符合 同個字母數量都能 mod k == 0
            if(remain < 0) break; //如果 < 0，表示要比題目原本字串的長度更長才能符合，因此退出
        }
        if(remain < 0) continue; //判斷下一個字元替換此 s[i]

        //如果還有剩下的長度可以被使用，那我們就全加給 'a' 就好，不需要擔心 a 字母會不會 % k != 0，
        //因為剩下的長度必定會 %k == 0，否則我們整理的第一點就不會符合。
        //證明: 每個字母的長度必須都 %k == 0，因此如果有剩下的長度 > 1，就表示題目的長度 %k != 0
        cnt[0] += remain;
        return i; //表示此 i+'a' 此字串替換給 s[i] 是好選擇，所以回傳
    }
    return 0;
}

string solve(){
    if(s.length() % k != 0) return "-1"; //我們整理的第一個重點

    for(int i = 0; i < 26; i++) sum[0][i] = 0; //歸零，不要用 memset，我這樣更省時間
    for(int i = 0; i < n; i++){ //開始計算在每個長度時，每個字元的數量
        sum[i][s[i] - 'a'] += 1;
        for(int j = 0; j < 26; j++) sum[i+1][j] = sum[i][j]; //狀態更新給下一個長度
    }

    int flag = 1; //判斷是否不需要修改
    for(int i = 0; i < 26; i++){
        if(sum[n-1][i] % k != 0){ //需要修改
            flag = 0; 
            break; //退出
        }
    }
    if(flag == 1) return s; //如果是 1 表示不需要修改，就是我們整理的第二個重點

    string ans = ""; //答案
    for(int i = n-1; i >= 0; i--){ //從字串最尾巴開始進行修改
        flag = check(i); //判斷這邊修改為哪個字元更好，0 表示都不好
        if(flag){ //如果 > 0，表示有可以修改的字元
            ans = s.substr(0, i) + (char) (flag+'a'); 
            //原本的字串長度從 0 到 i-1 加上修改的字元，
            //注意 C++ 的 substr(position, len) 函式與其他語言語法不同

            //透過字典序依序將每個字母需要的數量補入 string，這裡的 string(次數, 字元)
            for(int i = 0; i < 26; i++) ans += string(cnt[i], i+'a');
            return ans; //回傳答案
        }
    }
    return ans;
}

int main()
{
#ifdef LOCAL
    freopen("in1.txt", "r", stdin);
#endif // LOCAL
    cin >> T;
    while(T--){ //輸入題目資料
        cin >> n >> k >> s;
        cout << solve() << '\n'; //給 sovle 回傳答案
    }
    return 0;
}