G - Divide a Sequence Editorial by penguinman
\(\text{dp}_i\) を、\(A\) の先頭 \(i\) 項からなる数列 \(A'=(A_1,A_2,\ldots,A_i)\) について問題を解いた場合の答えと定義します。すると、以下のような漸化式が成立します(ここで \(\text{dp}_0=1\) とする)。
- \(\text{dp}_i=\sum_{j=0}^{i-1} \text{dp}_j \times (\max(\{A_{j+1},A_{j+2},\ldots,A_i\})-\min(\{A_{j+1},A_{j+2},\ldots,A_i\}))\)
このような動的計画法の計算量は愚直に行うと \(O(N^2)\) となりますが、工夫により \(O(N)\) まで落とすことが可能です。
まず、上記の漸化式を \(\max\) と \(\min\) それぞれについてで分離します。
- \(\text{dp}_i=\sum_{j=0}^{i-1} \text{dp}_j \times \max(\{A_{j+1},A_{j+2},\ldots,A_i\}) - \sum_{j=0}^{i-1} \text{dp}_j \times \min(\{A_{j+1},A_{j+2},\ldots,A_i\})\)
こうして得られる遷移の形、具体的には「\(\sum_{j=0}^{i-1} \text{dp}_j \times \max(\{A_{j+1},A_{j+2},\ldots,A_i\})\)」といった形は非常に高速化との相性が良いです。stack 等を用いて \(\max(\{A_{j+1},A_{j+2},\ldots,A_i\})\) の値を更新しつつ累積和を計算していくことで、線形時間で \(\text{dp}\) を行うことができます。
具体的には、stack に以下の \(2\) つをセットにして積んでいき、適切に差分を更新していけばよいです。
- \(\max(\{A_{j+1},A_{j+2},\ldots,A_i\})\) の値
- \(\max(\{A_{k+1},A_{k+2},\ldots,A_i\})\) の値が上に等しくなるようなすべての \(k\) についての \(\text{dp}_k\) の総和
\(\min\) パートについてもほぼ同様です。
実装例 (C++)
#include<bits/stdc++.h>
using std::cin;
using std::cout;
using std::endl;
using std::vector;
using ll = long long;
const ll MOD = 998244353;
int main(){
int N; cin >> N;
vector<ll> A(N);
for(int i=0; i<N; i++) cin >> A[i];
vector<ll> dp(N+1);
dp[0] = 1;
std::stack<ll> max, min, max_v, min_v;
ll max_sum = 0, min_sum = 0;
for(int i=0; i<N; i++){
// max
{
ll sum_v = dp[i];
while(!max.empty() && max.top() < A[i]){
max_sum -= max_v.top()*max.top();
sum_v += max_v.top();
max_sum %= MOD;
sum_v %= MOD;
max.pop();
max_v.pop();
}
max_sum += sum_v*A[i];
max_sum %= MOD;
max.push(A[i]);
max_v.push(sum_v);
}
// min
{
ll sum_v = dp[i];
while(!min.empty() && min.top() > A[i]){
min_sum -= min_v.top()*min.top();
sum_v += min_v.top();
min_sum %= MOD;
sum_v %= MOD;
min.pop();
min_v.pop();
}
min_sum += sum_v*A[i];
min_sum %= MOD;
min.push(A[i]);
min_v.push(sum_v);
}
dp[i+1] = max_sum-min_sum;
dp[i+1] %= MOD;
}
ll ans = dp[N];
if(ans < 0) ans += MOD;
cout << ans << endl;
}
なお、この他にも Cartesian Tree の構造に注目した \(\text{dp}\) 高速化手法も存在します(参考:https://atcoder.jp/contests/arc115/tasks/arc115_e )
posted:
last update: