2021牛客寒假算法基础集训营5 D.石子游戏

题目描述

叶妹妹很喜欢玩石头，于是这天泽鸽鸽给她出了一道石子游戏，规则是这样的：有 $n$ 堆石子排成一行，其中第 $i$ 堆石子有 $a_i$ 个，叶妹妹可以选择做无数次这种操作：每次操作把连续相邻的 $k$ 个石子堆中的每堆石子数目加一，请问叶妹妹能否让每堆石子的数目都相同呢？叶妹妹觉得这题太简单了，于是丢给了聪明的你，快来解决这个问题吧！

输入描述:

第一行输入样例组数 $T$

对于每组样例来说，第一行输入两个数 $n$ 和 $k$

第二行输入 $n$ 个数，其中第 $i$ 个数为 $a_i$

输出描述:

输出总共 $T$ 行，对于每组样例来说，如果能使每堆石子的数目都相同则输出一个整数 $x$，$x$ 表示达到相同时的最少的操作数；否则输出 $-1$

示例1

输入

1
2
3

1
4 3
1 1 1 2

输出

思维+差分~
这题我一开始想的是模拟，后来发现复杂度上天，忽然发现，可以用两遍差分来判断，我们先从左往右差分，在遍历的过程中记录最大值，并不断将小于最大值的数修正，不难发现，一遍差分之后该数列下标 $[1,n-k+1]$ 的部分（因为限定连续 $k$ 个数，所以右端点只能到 $n-k+1$）会变成非严格上升子序列，此时只需要再从右往左再进行一遍差分，判断数组是否修正成一样的数即可~
想法很简单，但是一边差分一边修正这个过程真的很难写，一不小心就挂了，我具体讲一下用到的两个变量：

diff 数组，记录所有位置的操作次数

sum，记录差分的前后缀和

这些变量的位置尤为重要，首先我们要先更新 sum，其次更新 a[i]，最后再更新 diff 数组，详见代码：

#include <bits/stdc++.h>

using namespace std;
typedef long long ll;
const int N = 1e5 + 5;
int t, n, k;
ll a[N], diff[N], mx, sum, ans;

void init() {
    mx = sum = 0;
    memset(diff, 0, sizeof(diff));
}

int main() {
    scanf("%d", &t);
    while (t--) {
        ans = 0;
        scanf("%d%d", &n, &k);
        for (ll i = 1; i <= n; i++) scanf("%lld", &a[i]);
        init();
        for (ll i = 1; i <= n; i++) {
            sum += diff[i];
            a[i] += sum;
            if (i <= n - k + 1 && mx > a[i]) {
                ans += mx - a[i], diff[i] += mx - a[i], diff[i + k] -= mx - a[i], sum += mx - a[i], a[i] = mx;
            }
            mx = max(mx, a[i]);
        }
        init();
        for (ll i = n; i >= 1; i--) {
            sum += diff[i];
            a[i] += sum;
            if (i >= k && mx > a[i]) {
                ans += mx - a[i], diff[i] += mx - a[i], diff[i - k] -= mx - a[i], sum += mx - a[i], a[i] = mx;
            }
            mx = max(mx, a[i]);
        }
        int flag = 0;
        for (ll i = 1; i <= n - 1; i++) {
            if (a[i] != a[i + 1]) {
                flag = 1;
                break;
            }
        }
        if (flag) printf("-1\n");
        else printf("%lld\n", ans);
    }
    return 0;
}