2023牛客暑期多校训练营2

K - Box

https://ac.nowcoder.com/acm/contest/57356/K

题意

给出一个长度为 $n$ 的 $01$ 串以及每个位置都有权值 $a_i \geq 0$ ，每个 $1$ 最多移动一次且最多移动一位。若一个位置有 $1$ 则可以获得权值，问最大可能的权值和。

$n \leq 10^6$。

题解

状态表示：

$a_i:=$ 权重数组

$b_i:=$ $01$ 数组

$dp_{ij}:=$ 考虑数列 $1\sim i$ 位，$j=0$ 代表此处的 $1$ 挪到上一位，$j=1$ 代表此处的 $1$ 不动，$j=2$ 代表此处的 $1$ 挪到下一位，此时的最大权值和为 $dp_{ij}$.

状态计算：

  • 若 $b_i=0$

    • $dp_{i,0}=dp_{i-1,0}$
    • $dp_{i,1}=dp_{i-1,1}$
    • $dp_{i,2}=dp_{i-1,2}$

  • 若 $b_i=1$

    • 如果与上一个 $1$ 的距离为 $1$：初始状态为 $11$（下面的箭头代表 $1$ 的初始位置）

      • $dp_{i,0}=dp_{i-1,0}+a_{i-1}$：代表情况为 $1\underset{\hat{\ }}1\underset{\hat{\ }}00$
      • $dp_{i,1}=\max\{dp_{i-1,0},dp_{i-1,1}\}+a_{i}$：代表情况为 $\max\{1\underset{\hat{\ }}0\underset{\hat{\ }}10,0\underset{\hat{\ }}1\underset{\hat{\ }}10\}$
      • $dp_{i,2}=\max\{dp_{i-1,0},dp_{i-1,1},dp_{i-1,2}\}+a_{i+1}$：代表情况为 $\max\{1\underset{\hat{\ }}0\underset{\hat{\ }}01,0\underset{\hat{\ }}1\underset{\hat{\ }}01,0\underset{\hat{\ }}0\underset{\hat{\ }}11\}$

    • 如果与上一个 $1$ 的距离为 $2$：初始状态为 $101$

      • $dp_{i,0}=\max\{dp_{i-1,0},dp_{i-1,1}\}+a_{i-1}$：代表情况为 $\max\{1\underset{\hat{\ }}01\underset{\hat{\ }}00,0\underset{\hat{\ }}11\underset{\hat{\ }}00\}$
      • $dp_{i,1}=\max\{dp_{i-1,0},dp_{i-1,1},dp_{i-1,2}\}+a_{i}$：代表情况为 $\max\{1\underset{\hat{\ }}00\underset{\hat{\ }}10,0\underset{\hat{\ }}10\underset{\hat{\ }}10,0\underset{\hat{\ }}01\underset{\hat{\ }}10\}$
      • $dp_{i,2}=\max\{dp_{i-1,0},dp_{i-1,1},dp_{i-1,2}\}+a_{i+1}$：代表情况为 $\max\{1\underset{\hat{\ }}00\underset{\hat{\ }}01,0\underset{\hat{\ }}10\underset{\hat{\ }}01,0\underset{\hat{\ }}01\underset{\hat{\ }}01\}$

    • 如果与上一个 $1$ 的距离为 $\geq3$：初始状态为 $10\dots01$

      • $dp_{i,0}=\max\{dp_{i-1,0},dp_{i-1,1},dp_{i-1,2}\}+a_{i-1}$：代表情况为 $\max\{1\underset{\hat{\ }}00\dots1\underset{\hat{\ }}00,0\underset{\hat{\ }}10\dots1\underset{\hat{\ }}00,0\underset{\hat{\ }}01\dots1\underset{\hat{\ }}00\}$
      • $dp_{i,1}=\max\{dp_{i-1,0},dp_{i-1,1},dp_{i-1,2}\}+a_{i}$：代表情况为 $\max\{1\underset{\hat{\ }}00\dots0\underset{\hat{\ }}10,0\underset{\hat{\ }}10\dots0\underset{\hat{\ }}10,0\underset{\hat{\ }}01\dots0\underset{\hat{\ }}10\}$
      • $dp_{i,2}=\max\{dp_{i-1,0},dp_{i-1,1},dp_{i-1,2}\}+a_{i+1}$：代表情况为 $\max\{1\underset{\hat{\ }}00\dots0\underset{\hat{\ }}01,0\underset{\hat{\ }}10\dots0\underset{\hat{\ }}01,0\underset{\hat{\ }}01\dots0\underset{\hat{\ }}01\}$

代码

#include <bits/stdc++.h>
#define endl '\n'
#define int long long

using namespace std;

void solve()
{
    int n;
    cin >> n;
    vector<int> a(n + 10), b(n + 10);
    for (int i = 1; i <= n; i++)
        cin >> a[i];
    for (int i = 1; i <= n; i++)
        cin >> b[i];
    vector<vector<int>> dp(n + 10, vector<int>(3));
    int last = -2;
    for (int i = 1; i <= n; i++)
    {
        if (b[i] == 1)
        {
            if (i - last == 1)
            {
                dp[i][0] = dp[i - 1][0] + a[i - 1];
                dp[i][1] = max({dp[i - 1][0], dp[i - 1][1]}) + a[i];
                dp[i][2] = max({dp[i - 1][0], dp[i - 1][1], dp[i - 1][2]}) + a[i + 1];
            }
            else if (i - last == 2)
            {
                dp[i][0] = max({dp[i - 1][0], dp[i - 1][1]}) + a[i - 1];
                dp[i][1] = max({dp[i - 1][0], dp[i - 1][1], dp[i - 1][2]}) + a[i];
                dp[i][2] = max({dp[i - 1][0], dp[i - 1][1], dp[i - 1][2]}) + a[i + 1];
            }
            else
            {
                dp[i][0] = max({dp[i - 1][0], dp[i - 1][1], dp[i - 1][2]}) + a[i - 1];
                dp[i][1] = max({dp[i - 1][0], dp[i - 1][1], dp[i - 1][2]}) + a[i];
                dp[i][2] = max({dp[i - 1][0], dp[i - 1][1], dp[i - 1][2]}) + a[i + 1];
            }
            last = i;
        }
        else
        {
            dp[i][0] = dp[i - 1][0];
            dp[i][1] = dp[i - 1][1];
            dp[i][2] = dp[i - 1][2];
        }
    }
    cout << max({dp[n][0], dp[n][1], dp[n][2]}) << endl;
}

signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    solve();
    return 0;
}