[백준] 23245 Similarity

문제 링크

http://icpc.me/23245

23245번: Similarity

문제 출처는 2021 icpc 온라인 예선이다.

풀이

(pi < pj < pk)&& (qi < qj < qk) 인 (i, j, k)를 구하면 되는 문제이다.

pi < pj < pk인 수를 쉽게 구하기 위해 (pi, qi)를 쌍으로 묶어 정렬을 한다. 그렇다면 p는 오름차순 정렬이 된 상태를 유지하기 때문에 i < j < k인 쌍을 아무거나 고른다면 항상 _pi <= pj <= pk_를 유지한다.

그렇다면 j를 고정한 상황을 고려하면 (pj < pi && q0..j-1중 qj 보다 작은놈의 개수) * (pi < pk && qj+1... n 중 qj 보다 큰 놈의 개수)를 계산한다면 (i, j, k) 쌍의 개수를 구할 수 있다. 이때 두 값은 세그먼트 트리를 이용하여 O(logN)에 구할 수 있기 때문에 시간 복잡도는 O(NlogN)이 된다.

예제를 통해 이해해보자.

2 5 9 5 1
1 4 5 3 2
-> 정렬 후
1 2 5 5 9 
2 1 3 4 5

아래의 수에서 j를 3로 고정을 한 (i, 3, k) 쌍의 개수를 구해보겠다.

qj == 3이기 때문에 (pi < pj && 0 <= i <= 2 && qi < 3)인 수는 2, 1 두 개가 있고, (pj < pk && 4 <= k <= n && qk > 3) 은 5 한 개가 있다. 따라서 2 * 1로 (i, 3, k) 2개의 쌍을 만들 수 있다. p에서 같은 수가 나오는 부분은 세그먼트 트리의 업데이트 순서를 잘 조작하여 처리할 수 있다.

코드

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define mid (s + e >> 1)
#define f first
#define s second

class SegTree {
private:
    const int N;
    vector<ll> tree;
public:
    void update(int idx, ll diff, int node, int s, int e) {
        if (idx < s || idx > e) return;
        tree[node] += diff;
        if (s == e) return;
        update(idx, diff, 2*node, s, mid);
        update(idx, diff, 2*node + 1, mid + 1, e);
    }
    SegTree(int N) : N(N) { tree.resize(4 * N); }
    void update(int idx, ll diff) { update(idx, diff, 1, 0, N); }

    ll query(int qs, int qe, int node, int s, int e) {
        if (qs <= s && e <= qe) return tree[node];
        if (qe < s || e < qs) return 0;
        return query(qs, qe, 2*node, s, mid) + query(qs, qe, 2*node + 1, mid + 1, e);
    }
    ll query(int qs, int qe) { return query(qs, qe, 1, 0, N); }
};


pair<int, int> arr[100005];
long long ans;
constexpr int INF = 1e6 + 6;

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr);
    int n; cin >> n;
    SegTree t1(INF), t2(INF);
    for (int i = 0; i < n; i++) cin >> arr[i].f;
    for (int i = 0; i < n; i++) cin >> arr[i].s, t2.update(arr[i].s, 1);
    sort(arr, arr + n);

    for (int i = 0; i < n; i++) {
        // p가 같은 수를 처리하기 위해 한번에 업데이트
        int s = i, e = i;
        while (e < n - 1 && arr[i].f == arr[e + 1].f) e++;
        for (int j = s; j <= e; j++) t2.update(arr[j].s, -1);
        for (int j = s; j <= e; j++) ans += t1.query(0, arr[j].s - 1) * t2.query(arr[j].s + 1, INF);
        for (int j = s; j <= e; j++) t1.update(arr[j].s, 1);
        i = e;
    }
    cout << ans;
}