招商银行2020FinTech精英训练营-研发赛道 C 修塔游戏

题目描述

小招正在玩一款修塔游戏：系统中有n座高塔，每座高塔由若干个高度相同的方块堆砌而成。修塔游戏的规则为：
（1）每次从最高塔的塔尖拿走一个方块
（2）每次在最低塔的塔尖堆砌一个方块
小招每次只能完成上述两个动作中的一个动作。游戏的目标是使n座高塔中至少有k座高塔的高度相同，请问小招最少需要多少次才能完成游戏。

输入描述:

输入共有2行，第一行为n和k(1≤k≤n≤200000 )，第二行为n座塔的高度组成的数组 a1, a2, ...an(1≤aj≤10000)。

输出描述:

输出值为最少需要多少次动作才能完成游戏。

示例1

输入
6 5
1 2 2 4 2 3

输出
3

思路

假设最终操作完成k座高塔高度为h,则在操作之前的高塔中，存在高度小于h、等于h、大于h三类，每一类数量分别是n_neg,n_zero,n_pos。若n_zero大于等于k,则不需要任何操作。其他情况下，根据k的大小，在高度小于h和大于h的高塔中进行操作最终满足要求。存在一下几种情况

• 只操作高度小于h的塔
• 只操作高度大于的塔
• 高度小于h的塔都需要操作，又可以分为
• • 先操作高度小于h的塔，再操作高度小于h的塔
• • 先操作高度小于h的塔，再操作高度小于h的塔

所有情况取最小值得到在这个h下的最小操作次数，h可以取得范围是0~最高塔的高度，然后在h的这个范围内，取操作数最小值。

代码

#include<iostream>
#include<algorithm>
#include<climits>
using namespace std;

int count_step(vector<int> &a_sorted, int k){
    //vector<int> a_neg, a_pos, a_zero;
    int n_neg = 0, n_pos = 0, n_zero = 0;
    int orig_step_lf = 0, orig_step_hf = 0;
    for(auto ele:a_sorted){
        if(ele < 0){
            //a_neg.push_back(abs(ele+1));
            orig_step_lf+=abs(ele+1);
            n_neg++;
        }
        else if(ele == 0){
            //a_zero.push_back(ele);
            n_zero++;
        }
        else{
            //a_pos.push_back(ele);
            orig_step_hf+=abs(ele-1);
            n_pos++;
        }
    }

    if(n_zero >= k)
        return 0;

    //Lower First
    int step_lf = 0;
    if(n_neg > 0){
        if(n_neg >= k - n_zero)
            step_lf = orig_step_lf + k - n_zero;
        else{
            step_lf = orig_step_lf + orig_step_hf + k - n_zero;
        }
    }
    else{
        step_lf = orig_step_hf + k - n_zero;
    }

    //Higer First
    int step_hf = 0;
    if(n_pos > 0){
        if(n_pos >= k - n_zero)
            step_hf = orig_step_hf + k - n_zero;
        else
            step_hf = orig_step_hf + orig_step_lf + k - n_zero;
    }
    else{
        step_hf = orig_step_lf + k - n_zero;
    }

    return min(step_hf, step_lf);
}

int main()
{
    int n,k;
    cin>>n>>k;
    vector<int> a(n);
    for(int i=0;i<n;i++){
        cin>>a[i];
    }
    int max = *max_element(a.begin(),a.end());
    int min_step = INT_MAX;
    for(int i=0;i<=max;i++){
        vector<int> dif(n);
        for(int j=0;j<n;j++)
            dif[j] = a[j]-i;
        //a_sorted = sorted(dif)
        int step = count_step(dif, k);
        min_step = min(min_step, step);
    }
    cout<<min_step;
    return 0;
}