Leetcode 第 294 场周赛 题解

以后周赛的水题就直接略过了,直入主题。

4. 巫师的总力量和

这题很快就有想法了,大概 10 分钟左右吧。然后推 单调栈debug 用了不少时间,导致场上没过这题。说明我目前的薄弱点仍然是 单调栈模拟细节

题意

给你一个数组 $s[N]$,计算 $\sum_{1 \le i \le j \le N} min(i,j) \times sum(i,j)$。其中 $min(i,j) \times sum(i,j)$ 称为 $s[i:j]$ 的力量和。

思路

仍然以 后缀 的方向考虑。假设以 $s[i-1]$ 结尾数组的 总力量和 为 $dp(i - 1)$,考虑以 $s[i]$ 结尾的情况。

  1. 如果 $s[i]$ 是 $s[1:i]$ 的最小值,那么

    $dp(i) = s[i] \times \sum_{1 \le j \le i} sum(j,i)$。

    $=s[i] \times \sum_{1 \le j \le i} (s_1[i] - s_1[j-1])$

    $=s[i] \times ((i + 1) s_1[i] - \sum_{1 \le j \le i}s_1[j-1])$

    $=s[i] \times ((i + 1) s_1[i] - s_2[i - 1])$

    其中 $s_1[N]$ 为 $s[N]$ 的前缀和, $s_2[N]$ 为 $s1[N]$ 的前缀和。

  2. 否则,必然存在一个最大的 $k$,满足 $k \lt i$ 且 $s[k] \lt s[i]$。

    • 在 $s[k+1:i]$ 中,$s[i]$ 仍然是最小值,对答案贡献为

      $s[i] \times \sum_{k+1 \le j \le i} sum(j,i)$

      $=s[i] \times \sum_{k+1 \le j \le i} (s_1[i] - s_1[j-1])$

      $=s[i] \times [(i - k)s_1[i] - (s_2[i - 1] - s_2[k - 1])]$

    • 在 $s[1:k]$ 中,$s[i]$ 不是最小值了,该区间对答案贡献为

      $\sum_{1\le j \le k} min(j,i) \times sum(j,i)$

      $=\sum_{1\le j \le k} min(j,k) \times [sum(j,k) + sum(k+1,i)]$

      $=\sum_{1\le j \le k} min(j,k) \times sum(j,k) + sum(k+1,i) \times \sum_{1\le j \le k} min(j,k)$

      $=dp(k) + (s_1[i] - s_1[k]) \times \sum_{1\le j \le k} min(j,k)$

      $=dp(k) + (s_1[i] - s_1[k]) \times dp_2(k)$

不妨设 $dp_2(i) = \sum_{1\le j \le i} min(j,i)$。

  1. 如果 $s[i]$ 是 $s[1:i]$ 的最小值,那么 $dp_2(i)=i \times s[i]$
  2. 否则,必然存在一个最大的 $k$,满足 $k \lt i$ 且 $s[k] \lt s[i]$
    • 那么在 $s[k+1:i]$ 中,$s[i]$ 仍然是最小值,对 $dp_2(i)$ 贡献为 $(i - k) \times s[i]$
    • 在 $s[1:k]$ 中,$s[i]$ 不是最小值了,该区间对 $dp_2$ 贡献为 $dp_2(k)$

现在 $dp(i)$ 的计算问题都解决了。问题只剩下如何寻找满足 $s[k] \lt s[i]$ 且 $k \lt i$ 的最大的 $k$。很明显这是一个 NGE (Next Greater Element) 问题,可以使用 单调栈 寻找最近的满足 $s[k] \lt s[i]$ 的元素。

因此该题是 单调栈优化动态规划 问题,时间复杂度 $O(N)$,空间复杂度 $O(N)$。

using ll = long long;
const int MOD = 1e9 + 7;

class Solution {

public:
    int totalStrength(vector<int> &strength) {
        int n = strength.size();
        auto sum1 = vector<ll>(n + 1);
        auto sum2 = vector<ll>(n + 1);

        auto dp1 = vector<ll>(n + 1);
        auto dp2 = vector<ll>(n + 1);


        ll ans = 0;
        deque<int> q;
        for (int i = 1; i <= n; ++i) {
            sum1[i] = (sum1[i - 1] + strength[i - 1]) % MOD;
            sum2[i] = (sum2[i - 1] + sum1[i]) % MOD;

            while (!q.empty() && strength[q.back()] > strength[i - 1]) {
                q.pop_back();
            }

            if (q.empty()) {
                dp2[i] = (1LL * strength[i - 1] * i) % MOD;
                dp1[i] = (dp1[i] + strength[i - 1] * sum1[i] % MOD * i % MOD) % MOD;
                dp1[i] = (dp1[i] - strength[i - 1] * sum2[i - 1] % MOD) % MOD;
            } else {
                int k = q.back() + 1;
                dp2[i] = (dp2[k] + 1LL * strength[i - 1] * (i - k)) % MOD;
                dp1[i] = (dp1[k] + (sum1[i] - sum1[k]) % MOD * dp2[k] % MOD) % MOD;

                dp1[i] = (dp1[i] + strength[i - 1] * sum1[i] % MOD * (i - k) % MOD) % MOD;
                dp1[i] = (dp1[i] - strength[i - 1] * (sum2[i - 1] - sum2[k - 1]) % MOD) % MOD;
            }

            dp1[i] = (dp1[i] + MOD) % MOD;
            q.push_back(i - 1);
            ans = (ans + dp1[i]) % MOD;
        }
        return ans;
    }
};

插眼

Leetcode 上看到更好,更清晰的做法。相比于使用 后缀 思想考虑,不如直接以 贡献 思想考虑。即考虑 $s[i]$ 影响的范围 和对 答案的贡献

  1. 为什么会想到这样做?不妨看看我们要求的式子:

    $\sum_{i=1}^{N}\sum_{j=i}^{N} min(i,j) \times sum(i,j)$

    $=\sum_{i=1}^{N}\sum_{j=i}^{N} s[argmin(i,j)] \times sum(i,j)$

  2. 我们枚举 $k=argmin(i,j)$。但 $argmin(i,j)$ 具有多值性。我们设 $argmin(i,j)$ 为 区间内第一个最小的元素,那么 $argmin(i,j)$ 的值是唯一的。于是可以另 $k=argmin(i,j)$,枚举 $k$

    $\sum_{k=1}^{N}\sum_{i=1}^{N}\sum_{j=i}^{N} s[k] \times (k = argmin(i,j))\times sum(i,j)$

  3. 由于 $i \le k \le j$, 那么 $(k = argmin(i,j)) = (k = argmin(i,k))\times(k = argmin(k,j))$

    $\sum_{k=1}^{N}\sum_{i=1}^{N}\sum_{j=i}^{N} s[k] \times (k = argmin(i,j))\times sum(i,j)$

    $=\sum_{k=1}^{N}s[k]\sum_{i=1}^{N}(k = argmin(i,k))\sum_{j=i}^{N}(k = argmin(k,j)) sum(i,j)$

    $=\sum_{k=1}^{N}s[k]\sum_{i=l_k}^{k}\sum_{j=k}^{r_k} sum(i,j)$

    $l_i$ 为 满足 $k = argmin(i,k)$ 的最小的 $i$,$r_i$ 为 满足 $k = argmin(k,j)$ 的最大的 $j$。

  4. 我们另 $s_1[N]$ 为 $s[N]$ 的前缀和,$s_2[N]$ 为 $s_1[N]$ 的前缀和,那么

    $\sum_{k=1}^{N}s[k]\sum_{i=l_k}^{k}\sum_{j=k}^{r_k} sum(i,j)$

    $=\sum_{k=1}^{N}s[k]\sum_{i=l_k}^{k}\sum_{j=k}^{r_k} (s_1[j] - s_1[i-1])$

    $=\sum_{k=1}^{N}s[k][(k - l_k + 1)(s_2[r_k] - s_2[k - 1]) - (r_k - k + 1)(s_2[k-1]-s_2[l_k-2])]$

由于 $k=argmin(i,j)$ 为 区间内第一个最小的元素,可使用 单调栈 计算 $l_i$ 和 $r_i$。时间、空间复杂度均为 $O(N)$。

启示

  • 使用 $min(i,j)=s[argmin(i,j)]$ 前提是,$argmin(i,j)$ 必须是单值的。
  • 降低运算复杂度的本质是去除依赖。文中最关键的一步是 3.,将计算复杂度大大降低。

  • 换枚举变量 2. 也很关键,相当于换了一个维度处理问题。
using ll = long long;
const int MOD = 1e9 + 7;

class Solution {

public:
    int totalStrength(vector<int> &strength) {
        int n = strength.size();
        auto l = vector<int>(n + 1);
        auto r = vector<int>(n + 1);

        auto s1 = vector<ll>(n + 1);
        auto s2 = vector<ll>(n + 1);

        auto dq = deque<int>();

        for (int i = 1; i <= n; ++i) {
            s1[i] = (s1[i - 1] + strength[i - 1]) % MOD;
            s2[i] = (s2[i - 1] + s1[i]) % MOD;

            while (!dq.empty() && strength[dq.back()] > strength[i - 1]) {
                dq.pop_back();
            }
            l[i] = dq.empty() ? 1 : dq.back() + 2;
            dq.push_back(i - 1);
        }
        dq.clear();

        for (int i = n; i >= 1; --i) {
            while (!dq.empty() && strength[dq.back()] >= strength[i - 1]) {
                dq.pop_back();
            }
            r[i] = dq.empty() ? n : dq.back();
            dq.push_back(i - 1);
        }

        auto sum = [&s2](int x) -> ll {
            if (x <= 0) return 0;
            if (x >= s2.size()) return s2.back();
            return s2[x];
        };

        ll ans = 0;
        for (int i = 1; i <= n; ++i) {
            ll s = (i - l[i] + 1) * (sum(r[i]) - sum(i - 1)) % MOD - (r[i] - i + 1) * (sum(i - 1) - sum(l[i] - 2)) % MOD;
            ans = (ans + s * strength[i - 1]) % MOD;
        }
        return (ans + MOD) % MOD;
    }
};