小木棍 |豆包MarsCode AI刷题

小S有n根木棍，每根木棍上放着一定数量的算珠，分别为a_1, a_2, ..., a_n。他可以最多进行k次操作，每次操作允许从一根木棍上取下一颗算珠并放到另一根木棍上。小S希望通过这些操作，使得最终的算珠数量数组按照字典序尽可能小。 字典序小意味着：如果a和b是两个数组，当存在第一个不同的元素时，a中该元素应小于b中的对应元素。 你需要帮助小S找到最终可以得到的最小字典序数组。

1. 对初始算珠数量数组进行排序：

   • 首先使用 Python 的内置函数 sorted 对输入的算珠数量数组 a 进行排序，得到一个新的有序数组 sorted_a。这么做的目的是后续方便按照从小到大的顺序来调整算珠数量，以使得最终数组字典序尽可能小。

2. 通过循环进行算珠移动操作：

   • 初始化变量 i 为 0，用于控制循环遍历排序后的数组 sorted_a 的索引，同时作为判断当前比较和操作的相邻两个元素的位置。进入循环，循环条件是操作次数 k 大于 0 且索引 i 小于 n - 1（n 为木棍数量，保证在处理相邻元素时不会越界），在循环中进行以下操作：

     • 计算当前相邻两个元素（索引为 i 和 i + 1）的差值 diff，即 diff = sorted_a[i + 1] - sorted_a[i]，这个差值表示要让前一个元素和后一个元素相等需要移动的算珠数量（如果要让字典序更小，通常是将后一个元素的算珠往前面的元素上移动）。
     • 确定实际可以移动的算珠数量 move_count，取 diff 和剩余操作次数 k 中的较小值（通过 min(diff, k) 实现），这是因为不能超过剩余可操作次数以及实际需要移动的算珠数量限制。
     • 根据确定的移动数量 move_count，对当前相邻的两个元素进行算珠移动操作，将前一个元素 sorted_a[i] 的算珠数量减去 move_count，同时将后一个元素 sorted_a[i + 1] 的算珠数量加上 move_count，模拟算珠从后一根木棍移动到前一根木棍的过程，即 sorted_a[i] -= move_count 和 sorted_a[i + 1] += move_count。
     • 将剩余操作次数 k 减去已经移动的算珠数量 move_count（k -= move_count），更新剩余可操作次数。
     • 检查经过算珠移动操作后，当前前一个元素 sorted_a[i] 是否变为 0，如果变为 0，说明这根木棍上已经没有算珠了，此时将索引 i 加 1，继续往后检查下一组相邻元素，看是否还能进行算珠移动操作来进一步减小字典序。

3. 返回最终结果：

   • 当循环结束后（要么操作次数 k 用完了，要么已经遍历完所有相邻元素能调整的情况），sorted_a 数组就是经过最多 k 次操作后得到的算珠数量数组，它按照字典序已经尽可能小了，将其返回即可。

时间复杂度方面，对数组 a 进行排序的操作时间复杂度通常为 （n 为数组长度，这里使用 Python 内置的排序算法，一般是基于比较的排序算法，时间复杂度符合这个量级），后续的循环操作中，每次循环最多执行常数级别的操作，循环最多执行 n - 1 次（取决于 k 的大小以及元素差值情况），所以循环部分的时间复杂度为 ，综合起来整体时间复杂度为 ，在处理常规数量的木棍以及操作次数时能够通过合理调整算珠分布来得到最小字典序的数组。