题目描述

难度分:$1900$

输入$n$、$k(1 \leq k \leq n \leq 3 \times 10^5)$和长为$n$的数组$a(-10^6 \leq a[i] \leq 10^6)$。

把$a$分割成恰好$k$个非空连续子数组。

设$S=$第一个子数组的元素和乘以$1+$第二个子数组的元素和乘以$2+…+$第$k$个子数组的元素和乘以$k$。

输出$S$的最大值。

输入样例$1$

5 2
-1 -2 5 -4 8

输出样例$1$

15

输入样例$2$

7 6
-3 0 -1 -2 -2 -4 -1

输出样例$2$

-45

输入样例$3$

4 1
3 -1 6 0

输出样例$3$

8

算法

贪心

让数组的索引从$0$开始，如果将一个长度为$4$的数组$a=[a_0,a_1,a_2,a_3]$划分为$[[a_0],[a_1,a_2],[a_3]]$，$S=a_0+(a_1+a_2)\times 2+a_3 \times 3$，本质上就是$(a_0+a_1+a_2+a_3)+(a_1+a_2+a_3)+a_3$，即$s[0]+s[1]+s[3]$，$s[i]$为后缀$[i,n)$的累加和。

这样就很容易看出思路了，$s[0]$是必选的，剩下$k-1$个后缀和选$s[1…n-1]$最大的$k-1$个即可。初始化答案$ans=s[0]$，然后把$s[1…n-1]$追加到一个数组$sums$中，排个序将最大的$k-1$个累加到$ans$上就能得到最终答案。

复杂度分析

时间复杂度

预处理后缀和时间复杂度为$O(n)$；对后缀和排序时间复杂度为$O(nlog_2n)$；把除了$s[0]$之外最大的$k-1$个后缀和加起来时间复杂度为$O(k)$。因此，时间复杂度为$O(nlog_2n)$。

空间复杂度

后缀和数组$s$的空间是线性的，为$O(n)$；将除了$s[0]$之外的后缀和追加到一个额外的$sums$数组中，也占用了$O(n)$的空间（这个空间可以省掉，但不影响空间复杂度，无关紧要）；排序的空间复杂度参考快排，为$O(log_2n)$。因此，整个算法的空间瓶颈就是线性的，额外空间复杂度为$O(n)$。

python 代码

n, k = map(int, input().split())
a = list(map(int, input().split()))
s = [0]*n
s[-1] = a[-1]
sums = [s[-1]]
for i in range(n - 2, -1, -1):
    s[i] = s[i + 1] + a[i]
    if i > 0:
        sums.append(s[i])
sums.sort(reverse=True)
ans = s[0] + sum(sums[:k - 1])
print(ans)