小红的数组切割

[牛客链接](https://www.nowcoder.com/practice/ddd09898956044a5bfa0995adf30ce29)

思路

贪心 + 前缀和

每个元素 $i$ 的权值为 $op(i) \times (a_i + j)$ ，其中 $op(i) = 1$ 若 $s_i = \text{'1'}$ ， $op(i) = -1$ 若 $s_i = \text{'0'}$ ， $j$ 为该元素所在块的编号。

将总权值拆成两部分：

$ $\sum_{i=1}^{n} op(i) \times (a_i + j_i) = \underbrace{\sum_{i=1}^{n} op(i) \times a_i}_{\text{固定项}} + \underbrace{\sum_{i=1}^{n} op(i) \times j_i}_{\text{受切割影响}}$ $

第一部分和切割无关，是个常数。关键在第二部分。

切割对块编号的影响：假设不切割，所有元素在第 1 块，此时第二部分的贡献为 $\sum op(i)$ 。每在位置 $p$ （第 $p$ 和 $p+1$ 个元素之间）添加一刀，位置 $p+1$ 及之后所有元素的块编号都加 1，总权值增量为：

$ $\Delta_p = \sum_{i=p+1}^{n} op(i)$ $

即从位置 $p+1$ 到末尾的 $op$ 后缀和。

我们最多切 $k-1$ 刀（得到 $k$ 块），因此只需从 $n-1$ 个候选切割位置的增量中，贪心选取最大的 $k-1$ 个正值累加即可。

以样例验证： $a=[1,2,3,4]$ ， $s=\text{"1001"}$ ， $k=2$ 。

$op = [1, -1, -1, 1]$
固定项 $= 1 - 2 - 3 + 4 = 0$
不切割基础值 $= 0 + (1 - 1 - 1 + 1) = 0$
三个切割位置的增量： $\Delta_1 = -1-1+1 = -1$ ， $\Delta_2 = -1+1 = 0$ ， $\Delta_3 = 1$
选 $k-1=1$ 个最大正值： $\Delta_3 = 1$
答案 $= 0 + 1 = 1$ ，正确

代码

C++
Java

#include <bits/stdc++.h>
using namespace std;

int main() {
    int n, k;
    scanf("%d%d", &n, &k);
    vector<long long> a(n);
    for (int i = 0; i < n; i++) scanf("%lld", &a[i]);
    char s[100001];
    scanf("%s", s);

    vector<int> op(n);
    for (int i = 0; i < n; i++) {
        op[i] = (s[i] == '1') ? 1 : -1;
    }

    // 基础值：所有元素在第 1 块
    long long base = 0;
    for (int i = 0; i < n; i++) {
        base += (long long)op[i] * a[i] + op[i];
    }

    // 计算每个切割位置的增量（op 的后缀和）
    vector<long long> gains;
    long long suffix = 0;
    for (int i = n - 1; i >= 1; i--) {
        suffix += op[i];
        gains.push_back(suffix);
    }

    sort(gains.rbegin(), gains.rend());

    long long ans = base;
    for (int i = 0; i < k - 1 && i < (int)gains.size(); i++) {
        if (gains[i] > 0) {
            ans += gains[i];
        } else {
            break;
        }
    }

    printf("%lld\n", ans);
    return 0;
}

import java.util.*;
import java.io.*;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int k = Integer.parseInt(st.nextToken());

        long[] a = new long[n];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < n; i++) {
            a[i] = Long.parseLong(st.nextToken());
        }

        String s = br.readLine().trim();
        int[] op = new int[n];
        for (int i = 0; i < n; i++) {
            op[i] = (s.charAt(i) == '1') ? 1 : -1;
        }

        // 基础值：所有元素在第 1 块
        long base = 0;
        for (int i = 0; i < n; i++) {
            base += (long) op[i] * a[i] + op[i];
        }

        // 计算每个切割位置的增量
        long[] gains = new long[n - 1];
        long suffix = 0;
        for (int i = n - 1; i >= 1; i--) {
            suffix += op[i];
            gains[i - 1] = suffix;
        }

        Arrays.sort(gains);

        long ans = base;
        int cnt = 0;
        for (int i = gains.length - 1; i >= 0 && cnt < k - 1; i--, cnt++) {
            if (gains[i] > 0) {
                ans += gains[i];
            } else {
                break;
            }
        }

        System.out.println(ans);
    }
}

复杂度分析

时间复杂度： $O(n \log n)$ ，瓶颈在排序增量数组。
空间复杂度： $O(n)$ ，存储数组和增量。