BOJ 13215 - Fish [C++]

문제

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

길이가 $N$인 수열 $A$와 정수 $K$가 주어지면, $1 \le l \le r \le N$이면서 $\sum \limits_{i=l}^{r} {A[i]} \ge K$인 $(l,r)$ 쌍의 수를 구하는 문제이다.

 

풀이

$l = 1$인 쌍의 수를 어떻게 구할 지부터 생각해보자.

$S[i] = S[i-1] + A[i]$    $(i > 1)$ 라고 정의하면, $S[i] >= K$인 $i$의 개수가 된다.

 

그럼, $l = 2$인 쌍의 수는 어떻게 구할까?

기존의 배열 $S$에서 $S[1]$을 삭제하고, 모든 원소에 $A[1]$씩 뺀 후 $K$ 이상인지 검사하면 된다.

거꾸로 생각하면, 기존의 배열 $S$에서 $S[1]$을 삭제하고, $K + A[1]$ 이상의 원소 개수를 세면 된다.

 

일반화하면, $l$이 고정되었을 때의 구하려는 쌍의 수는 배열 $S[l..N]$에서 $K + S[l-1]$ 이상의 원소 개수와 동일하다.

 

누적 합 배열 $S$를 구해놓은 뒤 중복 원소를 지원하는 pbds에 모두 넣어놓고, $1 \le l \le N$에 대해 order_of_key 함수를 이용하여 쌍의 수를 계산해 주자.

 

#include <iostream>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
using ll = long long;

#define ordered_set tree<ll, null_type, less_equal<ll>, rb_tree_tag, tree_order_statistics_node_update>
ordered_set os;
void os_erase(ll val) {
    auto it = os.find_by_order(os.order_of_key(val));
    if (*it == val) os.erase(it);
}

ll A[200001]{}, S[200001]{};
int main() {
    cin.tie(0)->sync_with_stdio(0);

    ll N, K;
    cin >> N >> K;

    for (int i = 1; i <= N; i++) cin >> A[i], S[i] = S[i - 1] + A[i], os.insert(S[i]);

    ll ans = 0;
    for (int i = 1; i <= N; i++) {
        ans += N - os.order_of_key(K) - (i-1);
        K += A[i];
        os_erase(S[i]);
    }
    cout << ans;

}