E - Jingle, Jingle, Presents!!

この問題は主に尺取り法を問います。

単に $[l,r]\ (1\le l\lt r\le N)$ を全探索してしまうと $O(N^2)$ となり、時間制限に間に合いません。では、どうしたらよいでしょうか？

問題の制約から、区間 $[l,r]$ に対して以下の性質があることが分かります。

• $A_l,A_{l+1},\dots ,A_r$ において、$r$ を増やすと任意の整数 $k$ が出現する個数が増加する
• $A_l,A_{l+1},\dots ,A_r$ において、$l$ を増やすと任意の整数 $k$ が出現する個数が減少する

よってすべての区間において単調性があることがうかがえます。
またこの問題では条件を満たす区間の数え上げを要求しているので、これは尺取り法というアルゴリズムを用いて実装することができます。

任意の整数 $k$ の個数カウントですが、制約上これは連想配列などを用いることが適切です。以上から次の手続きを行えばよいです。
便宜上、区間 $[l,r]$ に出現する任意の整数 $k$ の個数を $\text{cnt}(k)$ と表します。

• $l=1,2,\dots ,N$ に対し、

  • $r+1\le N$ かつ $\text{cnt}(A_{r+1})$ が $K$ 以下である限り、$r,\text{cnt}(A_{r+1})$ を $1$ ずつ増やすことを繰り返す

  • もし $r=l$ となるならば、$r$ を $1$ 増やす

  • $r\ne l$ ならば $l$ を $1$ だけ増やす準備をするため、$\text{cnt}(A_l)$を $1$ 減らす

明らかにこの手続きを踏むことにより、最初に出てくる $k$ の個数の上限を超えないかつ $0$ 以上であることが見て取れます。
厳密にいうと、以下のようなことが分かります。

• 任意の整数 $k$ に対し区間 $[1,N]$ における $k$ の出現回数 $\text{CNT}(k)$ とすると、区間 $[l,r]$ において $\text{cnt}(k)\ge 0,\ \text{cnt}(k)\le \text{CNT}(k)$ である

手続きを見ると $O(N^2)$ と全探索した時と同じ計算量に見えますが、実は $O(N)$ であることが分かっています。
(この計算量のヒミツについては、他の方々の記事などを参考にしてください。ここでは省略します)

以上からこの問題は $O(N\log_{}N)$ で実装することが可能です。

xxxxxxxxxx

#include <bits/stdc++.h>

#define rep(i,n) for(int i = 0; i < n; i++) //[0,n)

using namespace std;

using ll = long long;

​

int main(){

    int N,M; cin >> N >> M;

    vector<int> A(N); rep(i,N) cin >> A[i];

​

    map<int,int> mp;

    int r = 0;

    ll ans = 0;

​

    mp[A[0]]++;

    rep(l,N){

        while(r+1 < N && mp[A[r+1]]+1 <= M) mp[A[r+1]]++,r++;

        ans += r-l;

        if(l == r) r++;

        else mp[A[l]]--;

    }

    cout << ans;

}

​