けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Codeforces Manthan Codefest 18 (Div. 1 + Div. 2) F - Maximum Reduction (R2500)

考え方はすぐわかるけど実装が大変...

問題へのリンク

問題概要

長さ  n の数列  a_{0}, a_{1}, \dots, a_{n-1} と整数  k が与えられる。 この数列に対して次のような操作を繰り返す

  • 連続する k 個の区間についての最大値をすべて書き出す

この操作を 1 回行うごとに数列の長さは  k-1 だけ減ることになる。数列の長さが  k-1 以下になるまで繰り返したときの、各操作後の数列の値の総和の総和を 1000000007 で割ったあまりを求めよ。

制約

  •  2 \le k \le n \le 10^{6}

考えたこと

各操作後にどうなっているのかをシミュレートする系ではなく、数列の各要素がどの段階まで生き残り、結果何個分ずつカウントされるのかを考える系の問題...というのはすぐにわかった。

が、そこから詰め切るのがちょっと大変だね...こういうのを実家と言えるようになりたい。

#include <iostream>
#include <vector>
using namespace std;

const int MOD = 1000000007;
typedef pair<int,int> pint;
const int INF = 1<<30;

int n, k;
vector<long long> a;

long long solve() {
    // a[i] の前後にどこまで伸びるか
    vector<int> up(n), down(n);
    vector<pint> vec;
    vec.push_back(pint(INF, n));
    for (int i = n-1; i >= 0; --i) {
        while (vec.back().first < a[i]) vec.pop_back();
        up[i] = vec.back().second;
        vec.push_back(pint(a[i], i));
    }
    vec.clear();
    vec.push_back(pint(INF, -1));
    for (int i = 0; i < n; ++i) {
        while (vec.back().first <= a[i]) vec.pop_back();
        down[i] = vec.back().second;
        vec.push_back(pint(a[i], i));
    }

    // 階段
    vector<long long> height(n + 1, 0);
    int cur = 0;
    for (int i = 1; i <= n; ++i) {
        if ((i - 1) % (k - 1) == 0) ++cur;
        height[i] = height[i - 1] + cur;
    }

    // 各要素ごとにカウント
    long long res = 0;
    for (int i = 0; i < n; ++i) {
        long long cnt = height[up[i]-down[i]-1] - height[up[i]-i-1] - height[i-down[i]-1] - 1;
        cnt = (cnt % MOD + MOD) % MOD;
        res = (res + cnt * a[i] % MOD) % MOD;
    }
    return res;
}

int main() {
    scanf("%d %d", &n, &k);
    a.resize(n); for (int i = 0; i < n; ++i) scanf("%lld", &a[i]);
    cout << solve() << endl;
}