CCC 2022 S3 Good Samples

    1#include <bits/stdc++.h>
    2using namespace std;
    3typedef long long ll;
    4const long long INF = 1e17; //❄️
    5int main() {
    6    ios::sync_with_stdio(0); cin.tie(0); 
    7    int n, m;              
    8    ll k;
    9    cin >> n >> m >> k;
    10    m = min(n, m);
    11    vector<int> ans(n);
    12    int cur = 1;
    13    int index = 0;
    14    bool all = true;
    15    //first m distinct pitch
    16    for(int i = 0; i < m; i++){
    17        // this checks if we can still place it or not since we have to remain the slots for placing one -> each slot must place at least one
    18        if(n - (i + 1) + i + 1 > k){
    19            all = false;
    20            break;
    21        }
    22        ans[i] = i + 1;
    23        index++;
    24        k -= i + 1;
    25    } 
    26    // cout << k << endl;
    27 
    28    //cannot place anything 
    29    if(index == 0){
    30        cout << -1 << '\n';
    31        return 0;
    32    } 
    33    //reset to form cycle 
    34    while(index < n && k > 0){
    35        // check if we can place the current pitch without going over -(n - index - 1) to account for the ones that we must place at the end
    36        if(k - m - (n - index - 1) >= 0 && all){
    37            ans[index] = cur;
    38            index++;
    39            cur++;
    40            if(cur > m) cur = 1;// the ensure it is distinct we do 1 to m repeatedly
    41            //every time we place a distinct pitch we add m good samples
    42            k -= m;
    43            // cout << "run" << '\n';
    44        } else {
    45            // we cannot place distinct pitch anymore
    46            int remain = n - index - 1; // how many ones that we must place at the end 
    47            int need = k - remain;
    48            int place = ans[index - need];
    49            ans[index] = place;
    50            k -= need;
    51            index++;
    52            // cout << remain << " " << need << " " << place << endl;
    53            break;
    54        }
    55        // cout << k << endl;
    56    }
    57    // for(int i = 0; i < ans.size(); i++){
    58    //     cout << ans[i] << (i == ans.size() - 1 ? "\n" : " ");
    59    // }
    60     //place the ones 
    61    while(index < n){
    62        ans[index] = ans[index - 1];
    63        index++;
    64        k--;
    65    }
    66    // cout << k << endl;
    67    if(k != 0){
    68        cout << -1 << '\n';
    69        return 0;
    70    }
    71    for(int i = 0; i < ans.size(); i++){
    72        cout << ans[i] << " \n"[i == ans.size() - 1];
    73    }
    74    return 0;
    75}