[BOJ 19833] Equal Maximums

https://www.acmicpc.net/problem/19833

 

간단한 수열에서의 문제. 앳코더 스타일 같아서 마음에 들었고 충분히 대회에도 나올만 한 문제라고 느꼈다. 다만 코딩미스라던지 생각을 깔끔히 정리하지 못해 시간이 꽤나 걸렸다..

다른 사람들의 답안을 훑어보니, 나와 비슷하게 푼 사람도 있었고 스택(?)을 이용하여 푸는 방법도 있는 것 같았다. 성실한 사람이라면 둘 모두 알아보면 좋을 것 같다.

 

문제

수열$\{a_i\}$가 주어진다. 서로 겹치는 원소가 없는 두 구간의 최댓값이 같은 경우의 수(즉, 구간쌍의 개수)를 구하시오.

 

풀이

수열의 수들을 좌표압축하면 $[1, n]$의 숫자로만 이루어져있다고 봐도 무방하다. 이제 $a_k=x$인 $k$들을 순서대로 $v_{x,1}, v_{x,2}, \cdots $라고 하자. 두 구간의 최댓값이 $i$인 경우를 합하면 답이 되므로, 이를 각각 구해보자.

 

$i$보다 큰 수들의 인덱스을 알고 있다고 생각해보자. 원하는 경우의 수를 겹치지 않게 세기 위해 나는 다음과 같은 방법을 사용했다. 최댓값이 $i$인 두 구간 중 왼쪽 구간에서 가장 오른쪽에 나타나는 $i$가 $v_{i,j}$에 존재하는 경우의 수를 각 $j$에 대해 세는 것이다. 이렇게 세면 겹칠 일이 없고 모든 경우를 셀 수 있다.

 

 

여기서 몇 가지 정의를 하자. $i$보다 크고 $v_{i,j}$보다 뒤에 있는 수들 중 가장 인덱스가 작은 수의 인덱스를 $U_{i,j}$, $i$보다 크고 $v_{i,j}$보다 앞에 있는 수들 중 가장 인덱스가 큰 수의 인덱스를 $L_{i,j}$라 하자.

 

최댓값이 $i$인 두 구간 중 왼쪽 구간에서 가장 오른쪽에 나타나는 $i$가 $v_{i,j}$에 존재하는 경우의 수를 $cnt_{i,j}$라 하면 $cnt_{i,j}=(\min(v_{i,j+1}, U_{i,j}) - v_{i,j})(v_{i,j}-L_{i,j})$임을 알 수 있다. 왼쪽, 오른쪽 방향으로 어디까지 늘릴 수 있는지 확인해보라.

 

또, 최댓값이 $i$인 구간 중 오른쪽 구간에서 가장 왼쪽에 나타나는 $i$가 $v_{i,j}$에 존재하는 경우를 $R_{i,j}$라 하고, $R_{i,j}+R_{i,j+1}+\cdots$를 $sum_{i,j}$라 하자. $sum_{i,j}$는 결론적으로 최댓값이 $i$인 구간 중 가장 왼쪽의 $i$가 $v_{i,j}$보다 크거나 같은 인덱스에 있는 구간의 개수이다.

$R_{i,j}$는 $cnt_{i,j}$를 구할 때에서 왼쪽 오른쪽만 바꿔 생각해주면 쉽게 구할 수 있다.

 

이제 두 가지 케이스를 나누자.

 

case 1: $v_{i,j+1} > U_{i,j}$

이 케이스에서는 최댓값이 $i$인 구간 중 $v_{i,j}$를 포함하는 구간과 $v_{i,j+1},v_{i,j+2}, \cdots$중 하나를 포함하는 구간은 겹칠 수 없다.  따라서 ($U_{i,j}$보다 오른쪽의 구간 중 최댓값이 $i$인 구간의 수의 개수) × (최댓값이 $i$인 구간 중 $v_{i,j}$에 위치한 $i$를 가장 오른쪽에 위치한 $i$로 가지는 구간의 개수)만큼의 경우의 생긴다. 이를 간략히 쓰면, $sum_{i,j+1} cnt_{i,j}$와 같다.

 

case 2: $v_{i,j+1} < U_{i,j}$

이 케이스에서는 최댓값이 $i$인 구간 중 $v_{i,j}$를 포함하는 구간과 $v_{i,j+2}, v_{i,j+3}, \cdots$ 중 하나를 포함하는 구간은 겹칠 수 없다. 이러한 경우의 수는 $sum_{i,j+2} cnt_{i,j}$와 같다.

오른쪽 구간이 $v_{i,j+1}$를 포함하는 경우만 더해주면 된다. $v_{i,j}$와 $v_{i,j+1}$ 사이에서 적당히 경계를 정한다면 $v_{i,j}$을 포함하는 구간의 왼쪽 확장 여부와 $v_{i,j+1}$를 포함하는 구간의 오른쪽 확장 여부(구간을 어디까지 늘릴 수 있는지)는 독립이므로, 따로 고려해보자. 우선 $v_{i,j}$의 왼쪽 확장과 $v_{i,j+1}$의 오른쪽 확장을 고려해보면 총 $(v_{i,j}-L_{i,j})(U_{i,j}-v_{i,j+1})$만큼의 경우의 수가 있고, 둘 사이의 경계를 정할 때 하나의 끝을 정하고 나머지 경우를 더해보면 $1+2+\cdots+(v_{i,j+1}-v_{i,j})=(v_{i,j+1}-v_{i,j})(v_{i,j+1}-v_{i,j}+1)/2$만큼의 경우의 수가 생긴다. 이 둘을 곱해서 더해주면 이 경우의 수 또한 구할 수 있다.

 

이를 자~알 더해주면 답이 나온다. 사실 코드를 보는게 이해가 빠를 수도.

$L_{i,j}, U_{i,j}$는 $i$보다 큰 수들의 인덱스가 들어간 set을 이용하여 구해주면 된다.

 

여담으로 어떤 $i$에 대해 $i$의 개수를 $k$라 하면 $v_{i,0}=0$, $v_{i,k+1}=n+1$라 가정, set에 처음부터 $0$과 $n+1$을 넣고 시작한다면 코딩을 깔끔하게 할 수 있다. 이 풀이 한정

 

코드

#include <bits/stdc++.h>
#define fi first
#define se second
#define eb emplace_back
#define em emplace
#define all(v) v.begin(), v.end()
#define reset(x) memset(x, 0, sizeof(x))

using namespace std;
typedef long long ll;
typedef pair <int, int> pii;
typedef pair <ll, ll> pll;

const int MAX = 101010;
const int INF = 1e9;
const ll LINF = 1e18;
const ll mod = 1e9+7;

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr);
    
    int n;
    cin >> n;
    vector <int> a(n+2), num;
    vector <vector <ll> > v(n+1);
    for(int i = 1; i <= n; i++) {
        cin >> a[i];
        num.eb(a[i]);
        v[i].eb(0);
    }
    sort(all(num));
    for(int i = 1; i <= n; i++) {
        a[i] = lower_bound(all(num), a[i]) - num.begin() + 1;
        v[a[i]].eb(i);
    }
    for(int i = 1; i <= n; i++) {
        v[i].eb(n+1);
    }

    set <ll> st;
    st.em(0), st.em(n+1);

    ll ans = 0;
    for(int i = n; i >= 1; i--) {
        vector <ll> sum(v[i].size());
        int k = v[i].size() - 2;
        for(int j = k; j >= 1; j--) {
            auto it = st.lower_bound(v[i][j]);
            ll temp = (*it - v[i][j]) * (v[i][j] - max(*prev(it), v[i][j-1]));
            sum[j] = sum[j+1] + temp;
            sum[j] %= mod;
        }
        for(int j = 1; j < k; j++) {
            auto it = st.lower_bound(v[i][j]);
            ll cnt = ((min(*it, v[i][j+1]) - v[i][j]) * (v[i][j] - *prev(it))) % mod;
            if(*it < v[i][j+1]) {
                ans += sum[j+1] * cnt;
                ans %= mod;
            }
            else {
                ans += sum[j+2] * cnt;
                ans %= mod;
                ll base = ((v[i][j] - *prev(it)) * (*it - v[i][j+1])) % mod;
                ll bet = ((v[i][j+1] - v[i][j]) * (v[i][j+1] - v[i][j] + 1) / 2) % mod;
                ans += (base * bet) % mod;
                ans %= mod;
            }
        }

        for(int j = 1; j <= k; j++) {
            st.em(v[i][j]);
        }

    }
    cout << ans;
}