题目描述

小招正在玩一款修塔游戏:系统中有n座高塔,每座高塔由若干个高度相同的方块堆砌而成。修塔游戏的规则为:

  1. 每次从最高塔的塔尖拿走一个方块
  2. 每次在最低塔的塔尖堆砌一个方块

小招每次只能完成上述两个动作中的一个动作。游戏的目标是使n座高塔中至少有k座高塔的高度相同,请问小招最少需要多少次才能完成游戏。


分析:

假设经过调整后,最终n座塔中k座塔的高度为h。数组a记录了n座塔的初始高度。算法对h进行遍历。

首先将塔分为三类,即高度高于h、低于h、等于h的,并记录每一类中塔的数量,分别为n_highern_lowern_equal。若高度等于h的塔有k座,即n_equal=k,则无需调整,步数为0。否则,可以有两种调整方案,分别为“先低后高”和“先高后低”。

以先低后高为例进行说明:该调整原则是指至少有一个高度低于h的塔在调整后高度变为h。由于只能对高度最低的塔进行操作,故在最后一步之前,所有高度低于h的塔的高度均应为h-1,否则将无法将其高度变为h,故在最后一步之前,必须将所有高度低于h的塔高调整为h-1,这需要的步长为数组a中所有小于h的数字与h-1的距离之和。此时,对于这些高度为h-1的塔,只需1步便可将其调整为h。若n_lower<=k-n_equal,则只需再从这n_lower个塔中,调整出k-n_equal个高度为h的塔即可,此时步长再增加k-n_equal。否则,若n_lower>k-n_equal,则还需在高度高于h的塔中调整出k-n_equal-n_lower个塔,调整方式与之前相同,即先将所有高度高于h的塔调整至h-1高度,再从其中调整k-n_equal-n_lower个塔即可。

先高后低的调整方式与先低后高类似。分别计算出两种方式需要的步数,选择其中最小值,即为将n座塔中k座塔的高度调整为h所需的最小步长。对于不同的h,再取极小值即可。

代码[Python]

# Count Step
def count_step(a_sorted, k):
    a_neg = []
    a_pos = []
    a_zero = []
    n_neg = n_pos = n_zero = 0
    for ele in a_sorted:
        if ele < 0:
            a_neg.append(ele)
            n_neg += 1
        elif ele == 0:
            a_zero.append(ele)
            n_zero += 1
        else:
            a_pos.append(ele)
            n_pos += 1

    if n_zero == k:
        return 0

    # Lower First
    step_lf = 0
    if n_neg > 0:
        step_lf = sum(list(abs(a_i + 1) for a_i in a_neg))
        if n_neg >= k - n_zero:
            step_lf = step_lf + k - n_zero
        else:
            step_lf = step_lf + n_neg
            step_lf = step_lf + sum(list(abs(a_i - 1) for a_i in a_pos))
            step_lf = step_lf + k - n_zero - n_neg
    else:
        step_lf = sum(list(abs(a_i - 1) for a_i in a_pos))
        step_lf = step_lf + k - n_zero - n_neg

    # Higer First
    step_hf = 0
    if n_pos > 0:
        step_hf = sum(list(abs(a_i - 1) for a_i in a_pos))
        if n_pos >= k - n_zero:
            step_hf = step_hf + k - n_zero
        else:
            step_hf = step_hf + n_pos
            step_hf = step_hf + sum(list(abs(a_i + 1) for a_i in a_neg))
            step_hf = step_hf + k - n_zero - n_pos
    else:
        step_hf = sum(list(abs(a_i + 1) for a_i in a_neg))
        step_hf = step_hf + k - n_zero - n_pos

    return min(step_hf, step_lf)

# Main Program
s = input().split(' ')
n = eval(s[0])
k = eval(s[1])
ls = input().split(' ')
a = list(eval(i) for i in ls)
min_step = 0

for i in range(max(a) + 1):
    dif = list((a_i - i) for a_i in a)
    a_sorted = sorted(dif)
    step = count_step(a_sorted, k)
    min_step = (step if i == 0 or step < min_step else min_step)
print(min_step)