[백준(BOJ)] 공장(7578번)_C++

문제를 보자마자 막연히 드는 풀이 방법은 A에서의 관계와 B에서의 관계가 다른 쌍을 찾는 방법이다.

즉 A에서는 i보다 j가 뒤에 있지만, B에서는 앞에있거나 혹은 그 반대의 관계이면 서로 교차한다.

그럼 모든 A에 대해 B를 검사하면O(n^2)의 시간 복잡도로 풀 수 있지만, n이 최대 50만이므로 시간 초과가 발생한다.

 

일단 교차하는지 확인해야 하기 때문에 모든 A 혹은 모든 B를 검사하긴 해야할 것 같다.

그럼 대충 O(nlogn)인 코드를 작성하도록 해보자.

 

모든 A에 대해 검사한다고 가정하자.

A[0]부터 먼저 처리하면 A[0]에 해당하는 B의 뒤에 있는 기계 중 이미 처리한 기계의 개수가 A[0]와 교차하는 기계의 개수다.

예제에서 132를 먼저 처리해보도록 하자. B에서 132 뒤에 있는 기계 중 먼저 처리된 기계가 있나? 없다. 132는 처리됐다.

392를 처리해보자. B에서 392 뒤에 있는 기계 중 처리된 기계가 있나? 132, 1개 존재한다. 그러면 392와 132는 교차한다.

B에서 311 뒤에 있는 기계 중 처리된 기계는 없다.

351 뒤에 있는 기계 중 처리된 기계는 132, 311, 2개 존재한다. 351과 132, 311은 교차한다.

231 뒤에 있는 기계는 없다.

위의 과정에서 우리가 찾은 교차 횟수는 3번이다.

 

그럼 대충 처리를 표시하는 bool 형 1차원 배열을 만들어서 처리하면 되겠다는 생각이 들 것이다.

그런데 이런식으로 구현하면 모든 A를 검사할 때마다 B에서 뒤에 있는 기계를 모두 검사해야 하므로 이또한 O(n^2)이다.

조금만 생각을 달리 해보면 B에서 뒤에 있는 기계를 모두 검사하는 이유는 처리된 기계들의 수를 세기 위해서다.

A에 따라 검사를 시작하는 위치가 계속 바뀌므로 구간합을 떠올릴 수 있다.

이 구간합을 세그먼트 트리를 이용해 구현하면 O(logn)만에 구간합을 구할 수 있고 최종적으로 O(nlogn)이 된다.

 

내가 구현한 코드를 간략하게 설명하자면 다음과 같다.

우선 A를 입력받고, B를 입력받는데 B의 경우에는 해당 기계에 대한 순서를 표시할 것이다.

그래야만 A에서의 기계에 대한 B에서의 위치를 O(1)로 구할 수 있다.

(B도 그대로 입력받으면 A에 해당하는 기계를 B에서 찾는 데에만 O(n)이 또 소요된다.)

 

세그먼트 트리는 어떤 의미를 갖는지 설명하자면

i가 리프노드에 해당하는 인덱스인 경우에는 tree[i]는 B[A[i]] 를 처리했는지 안했는지를 1과 0으로 표시하고,

i가 리프노드가 아닌 인덱스인 경우에는 i의 자식 노드들의 합이다. 즉 i의 자식 노드에 해당하는 기계들 중 처리된 기계들의 개수이다.

이러한 세그먼트 트리를 가지고 모든 A에 대해 검사하면서 처리 표시해주고, 트리 갱신하고를 반복하면 된다.

 

사실 글만 보고서 이해하기에는 좀 어려운 문제같다.

그래서 코드를 따라가면서 예제를 가지고 A, B배열과 세그먼트 트리를 직접 그려보면 생각보다 쉽게 이해할 수 있을 것이다.

 

 

---

정답 코드

#include <iostream>
#define ll long long

using namespace std;

int n, A[500'000], B[1'000'001];
ll tree[2'000'000], answer;

//주어진 구간의 합 구하기
ll sum(int cur, int start, int end, int sum_start, int sum_end){
    if (!(sum_start <= end && sum_end >= start))
        return 0;

    if (sum_start <= start && end <= sum_end)
        return tree[cur];

    int mid = (start + end) / 2;

    return sum(cur * 2, start, mid, sum_start, sum_end) +
            sum(cur * 2 + 1, mid + 1, end, sum_start, sum_end);
}


//방문한 노드 방문 처리하고, 트리 전체 갱신
void update(int cur, int start, int end, int visit_idx){

    if (!(visit_idx >= start && visit_idx <= end))
        return;
    
    if (start == end){
        tree[cur] = 1;
        return;
    }

    int mid = (start + end) / 2;

    update(cur * 2, start, mid, visit_idx);
    update(cur * 2 + 1, mid + 1, end, visit_idx);

    tree[cur] = tree[cur * 2] + tree[cur * 2 + 1];
}

int main(){
    cin.tie(NULL);
    cout.tie(NULL);
    ios_base::sync_with_stdio(false);

    //input
    cin >> n;
    for(int i = 0; i < n; i++)
        cin >> A[i];
    for(int i = 0; i < n; i++){
        int temp;
        cin >> temp;
        B[temp] = i;
    }

    //make segment tree
    for(int i = 0; i < n; i++){
        answer += sum(1, 0, n - 1, B[A[i]] + 1, n - 1);
        update(1, 0, n - 1, B[A[i]]);
    }

    //output
    cout << answer;
}