백준 - 공장 (7578) [C++]

이 문제는 세그먼트 트리를 이용해서 풀었다.

 

N <= 500,000이므로 O(N log N) 이하의 알고리즘을 구현해야한다.

 

이 문제에서 핵심은 A 배열을 앞에서부터 루프를 돌면서, 해당 A배열의 원소와 이어지는 B 배열의 인덱스보다 오른쪽에 연결된 것은 모두 교차한다는 사실을 찾아내는 것이다.

이를 찾은 이후는 간단하게 풀 수 있다. segment의 리프 노드의 경우 해당 index가 연결되었는지(B 배열 기준으로) 여부를 저장하게 되고, 구간에 대해서는 해당 구간에 연결된 노드의 수를 리턴하게 된다.

 

세그먼트 트리 문제는 규칙성을 찾아 세그먼트 트리로 풀어내는 능력이 중요한 것 같다.

#include <iostream>
#include <vector>
#include <algorithm>
#include <utility>
#include <string>
#include <map>
using namespace std;

#define MAX 500000

typedef long long ll;

int N;
int A[MAX];
int B[1000000 + 1];
ll segment[MAX * 4];

ll query(int node, int start, int end, int left, int right) {
	if (left > right) {
		return 0;
	}

	if (end < left || right < start) {
		return 0;
	}

	if (left <= start && end <= right) {
		return segment[node];
	}

	int mid = (start + end) / 2;
	return query(node * 2, start, mid, left, right) +
		query(node * 2 + 1, mid + 1, end, left, right);
}

void update(int node, int start, int end, int index) {
	if (index < start || end < index) {
		return;
	}

	segment[node]++;
	
	if (start != end) {
		int mid = (start + end) / 2;
		update(node * 2, start, mid, index);
		update(node * 2 + 1, mid + 1, end, index);
	}
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);

	cin >> N;
	for (int i = 0; i < N; i++) {
		cin >> A[i];
	}

	for (int i = 0; i < N; i++) {
		int tmp;
		cin >> tmp;

		B[tmp] = i;
	}
	
	ll answer = 0;
	for (int i = 0; i < N; i++) {
		int upNum = A[i];
		int bottomIndex = B[upNum];

		answer += query(1, 0, N - 1, bottomIndex + 1, N - 1);
		update(1, 0, N - 1, bottomIndex);
	}

	cout << answer << '\n';

	return 0;
}