ABC140 E - Second Sum (500)

E - Second Sum

概要

1からNまでの数字の順列Pが与えられる。1≦ L < R ≦ NなるすべてのL, Rの組について、P[L], P[L+1], ..., P[R]のうち2番目に大きい項を足していくことを考える。この総和を求めよ。

制約

2 ≦ N ≦ 10^5

方針

各L, Rについて2番目に大きい項を逐一探しに行くと計算が間に合わない。

あるP[i]が何回足されるかを考える。
P[i]が足されるときのL, i, Rの関係としては、P[L]からP[R]の間でP[i]より大きい項がちょうど1回現れる。これは次の2つの事象に分割して考えることができる。

• LからiまでにP[i]より大きい項がちょうど1回現れ、i+1からRまでにP[i]より大きい項が現れない
• LからiまでにP[i]より大きい項が現れず、i+1からRまでにP[i]より大きい項がちょうど1回現れる

よってiの右側、左側ごとにP[i]より大きい項が初めて現れるインデックスとP[i]が2回目に現れるインデックスを取得できれば、P[i]が何回足されるかを効率的に求めることが出来る。

インデックスの管理はmultisetを使い、P[i]が大きい方から作業していくとうまく実装できる。具体的には下記の通り。

• 0を2つ、N+1を2つmultisetにinsertしておく（番兵、計算の効率化）
• P[i]が大きい方からmultisetをlower_boundし、イテレータを動かすことで上記インデックスを取得する。なお、このときmultisetに含まれているインデックスは自分よりP[i]が大きいもののインデックスと番兵のみである。
• iをmultisetにinsertし、次のiでの作業を行う。

感想

公式解説、kmjpさんのブログを参考にしました。
kmjp.hatenablog.jp

解答

Submission #7580445 - AtCoder Beginner Contest 140

#include <bits/stdc++.h>
#define rep(i,n) for(int i=(0);i<(n);i++)
 
using namespace std;
 
typedef long long ll;
typedef pair<int, int> pii;
 
int main()
{
	cin.tie(0);
	ios::sync_with_stdio(false);
 
	int n;
	cin >> n;
 
	vector<int> p(n+1);
	rep(i, n) cin >> p[i+1];
 
	vector<pii> v(n);
	rep(i, n) v[i] = (pii){p[i+1], i+1};
	sort(v.begin(), v.end());
	reverse(v.begin(), v.end());
 
	multiset<int> ms;
	ms.insert(0);
	ms.insert(0);
	ms.insert(n+1);
	ms.insert(n+1);
 
	ll ans = 0;
	rep(i, n){
		int num = v[i].first;
		int id = v[i].second;
 
		auto itr = ms.lower_bound(id);
		ll y = *itr;
 
		itr++;
		ll z = *itr;
 
		itr--;
		itr--;
		ll x = *itr;
 
		itr--;
		ll w = *itr;
 
		ans += (x - w) * (y - id) * num;
		ans += (id - x) * (z - y) * num;
 
		// cout << w << " " << x << " " << id << " " << y << " " << z << endl;
		// cout << "add:" << (x - w) * (y - id) * num << endl;
		// cout << "add:" << (id - x) * (z - y) * num << endl;
		ms.insert(id);
	}
 
	cout << ans << endl;
}