[P5] 백준 7578번 공장 C++ 세그먼트 트리

리뷰

 

서로 교차하는 케이블 쌍의 개수를 구하고 출력하는 문제

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

 

전역 변수

• MAX_N : 입력 되는 노드의 최대 개수를 의미하는 정수형 상수 변수 500001로 초기화 해준다.
• nodes : 입력 받은 노드의 정보를 저장할 정수형 배열, MAX_N크기로 설정한다.
• index : 교차하기 위한 노드 정보를 저장할 정수형 배열, index를 값으로, value를 순서로 저장한다, MAX_N * 2크기
• tree : 세그먼트 트리 정보를 저장할 정수형 배열, MAX_N * 4크기로 설정한다.
• n, ans : 노드의 개수를 저장할 변수 n, 정답을 저장하고 출력할 변수 ans

 

함수

1. update

void update(ll node, ll s, ll e, ll idx)

 

1. 세그먼트 트리 정보를 업데이트 할 함수
2. 케이블이 연결되면 연결된 인덱스를 1로 바꿔준 후 구간합을 업데이트 해준다.

2. query

ll query(ll node, ll s, ll e, ll L, ll R)

 

1. L부터 R가지의 구간합을 구하는 함수
2. 구간에 포함되면 세그먼트 트리의 현재 노드를 출력한다.
3. 만약 포함되지 않는 부분이 있다면 mid를 기준으로 재귀적으로 이진탐색을 진행해 준다.

 

문제풀이

1. n값을 입력 받고 n개의 노드 정보를 nodes배열에 입력 받아준다.
2. n개의 숫자로 추가로 입력 받고, index의 인덱스로 설정한 뒤 현재 i번째 순번을 값으로 저장한다.
3. 다시 nodes배열을 탐색하며 해당 숫자를 index배열의 인덱스로 사용하여 인덱스를 찾아준다.
4. query 함수를 통해 해당 인덱스의 뒤 구간의 구간합을 모두 구해준 후에 ans에 더해준다.
5. update 함수를 통해 해당 인덱스의 값을 1로 바꿔주고 구간합을 업데이트 해준다.
6. 최종적으로 ans에 저장된 값을 출력해 준다.

 

참고 사항

최악의 경우 ans는 int범위를 벗어날 수 있다, long long타입으로 초기화 해준다.

 

A의 0번 인덱스 132를 처리해 줘야할 경우 index배열의 132번째 인덱스는 2가 될것이다.

그렇다면 query함수를 통해 3 ~ n - 1번째의 구간합을 구한 뒤 ans에 더해준다.

이후 update함수를 통해 세그먼트 트리에서 2번 노드에 해당하는 값을 1로 바꿔준다.

중복된 값이 나오지 않으므로 특정 노드가 1이상이 되는 경우는 없다.

 

정답 코드

#include<iostream>
#define ll long long

using namespace std;
const ll MAX_N = 500001;
ll nodes[MAX_N];
ll index[MAX_N * 2];
ll tree[MAX_N * 4];
ll n, ans;

void update(ll node, ll s, ll e, ll idx) {
	if (s == idx && idx == e) {
		tree[node] = 1;
		return;
	}
	ll mid = (s + e) / 2;
	if (s <= idx && idx <= mid) update(node * 2, s, mid, idx);
	else update(node * 2 + 1, mid + 1, e, idx);
	tree[node] = tree[node * 2] + tree[node * 2 + 1];
}

ll query(ll node, ll s, ll e, ll L, ll R) {
	if (R < s || e < L) return 0;
	if (L <= s && e <= R) return tree[node];
	ll mid = (s + e) / 2;
	ll left = query(node * 2, s, mid, L, R);
	ll right = query(node * 2 + 1, mid + 1, e, L, R);
	return left + right;
}

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

	cin >> n;
	for (int i = 0; i < n; i++) cin >> nodes[i];
	for (int i = 0; i < n; i++) {
		ll num; cin >> num;
		index[num] = i;
	}
	for (int i = 0; i < n; i++) {
		ll idx = index[nodes[i]];
		ans += query(1, 0, n - 1, idx + 1, n - 1);
		update(1, 0, n - 1, idx);
	}
	cout << ans;
}