[c++] 백준 - 공장 (인덱스드 트리)

1. 문제

2. 입출력 및 예제

3. 문제 해설

어떠한 힌트도 받지 않고 푼 첫번째 플레티넘 문제였다.

 

개인적으로는 백준-달리기(2517) 문제와 상당히 유사하다. 

 

문제를 이해하면 다음과 같다. 어떤 N개의 물체들이 특정한 번호를 가지면서 존재하는데 그 밑엔 다시 N개의 물체가 동일한 수를 가지며 순서가 섞여서 존재한다.

기존 N개에 써져있던 수가 동일한 물체끼리 연결을 하는데 이 때 교차한 쌍의 수를 구하면 된다. 

 

단순히 생각해보면 원래 물체에 써져있던 수는 하나도 중요하지 않다. 중요한 것은 그 물체가 원래 몇 번째에 존재했냐는 것이다. 

 

위의 예제를 생각해보면 다음과 같다. 

132 392 311 351 231

392 351 132 311 231

다음과 같은 물체가 존재할 때 꼬인 쌍의 수는 3이된다. 어떻게 나올까?

 

• 132는 1행에서 1번째에 나왔지만 2행에선 3번째에 나온다. 
  즉 1행에선 132보다 뒤에 있었지만 2행에서 132보다 앞에 존재하는 것이 2개(392 351)가 되는 것이다. 그러면 꼬인 쌍이 2개가 추가된다.
• 392는 1행에서 2번째에 나왔지만 2행에선 1번째에 나온다. 2행의 392보다 앞선 것은 없으므로 0개가 추가된다.
• 311은 1행에서 3번쨰에 나왔지만 2행에선 4번째에 나온다. 1행에서 311보다 뒤에 있던 것 중 2행에서 311보다 앞선 것은 351 1개이다. 따라서 꼬인 쌍에 1개가 추가된다. 
• 231은 1행에서 마지막에 나왔고 2행에서도 마지막에 나온다. 따라서 꼬인 쌍에 0개가 추가된다.  

위의 과정을 모두 거치면 총 3개의 꼬인 쌍이 나오게 된다. 

 

따라서 1행에서의 인덱스와 2행의 인덱스를 저장한 후 

1행에서 자기 인덱스보다 뒤에있던 것들이 2행에서 자기보다 앞서있는 쌍의 개수를 구해주면된다. 

---

트리를 구성할 데이터 준비

 

백준 - 달리기에서의 구성과 동일하다고 생각한다. 왜냐하면 이 문제는 결국 자신의 1행 인덱스보다 나중에 나왔던 것들이 2행에서 먼저 나오는 경우의 수를 구해주면 되기 때문이다. 

 

우선 앞서 말한 것을 구하려면 1행에서의 자신의 인덱스 2행에서 자신의 인덱스가 필요하다. 따라서 구조체를 선언해주었다. 

struct BB {
	int idxA;
    int idxB;
    bool inline const operator < (const BB &ref) const {
    	if(this->idxA > ref.idxA) return true;
        else return false;
    }
}

연산자 오버로딩은 추후 정렬을 위한 것이다. 

---

데이터 넣기

사실 위 문제에서 노드가 어떤 번호로 붙었는 지는 중요하지 않다. 따라서 idxA와 idxB만을 저장해주었다. 

이 때 2행에서 번호가 들어오면 그 번호에 해당하는 1행에서의 인덱스가 무엇인지 빠르게 접근하기 위해 hash_map을 기반으로 동작하는 unordered_map을 이용해 O(1)의 시간만에 인덱싱을 수행해주었다. 

int N; // 개수
int originA[500002]; // A를 입력받을 배열
int originB[500002]; // B를 입력받을 배열
unordered_map<int, int> mp; // A의 번호와 인덱스를 mapping할 map
void input() {
	scanf("%d", &N);
    for(int i=0; i <= N; i++) {
    	scanf("%d" , &originA[i]);
        mp[originA[i]] = i; // 인덱스와 번호 mapping
    }
    for(int i=0;  < N; i++) {
    	int tmp;
        scanf("%d" , &tmp);
        originB[i] = mp[tmp]; // mapping된 인덱스를 저장한다.
    }
}

---

인덱스드 트리 구성

인덱스드 트리를 구성하기 위해서 원본 배열을 저장할 offset이 필요하다. 이 offset은 N보다 큰 최초의 2의 제곱수에 해당한다. 

또한 트리를 구성해놓은 다음 답을 찾는 것이 아니라 데이터를 정렬한 후 넣으면서 트리의 query를 해야한다. 기본적인 인덱스드 트리의 query, update함수를 작성한다. 

int offset = 1;
int query(int from, int to) {
	int ans = 0;
	from += offset; to += offset;
    while(from <= to) {
    	if(from % 2) ans += Tree[from ++];
        if((to % 2) == 0) ans += Tree[to--];
        from /= 2; to /= 2;
    }
    return ans;
}

void update(int idx){
	idx += offset;
    Tree[idx] = 1; // 개수를 세는 것이므로 1만 추가한다.
    while(idx > 0) {
    	idx /=2;
        Tree[idx] = Tree[idx * 2] + Tree[idx * 2 + 1];
    }
}

int main() {
	input();
    while(offset < N)offset *=2; // 원본 배열을 저장할 위치를 만들어준다. 
}

---

핵심 아이디어

먼저 입력받은 2행의 번호들 (originB배열) 을 1행의 인덱스를 이용해서 내림차순으로 정렬해준다. 

왜냐하면 자기보다 1행에서 큰 인덱스를 가지는 것이 2행에서 더 앞에 나오는 지를 확인하기 위해서이다. 

 

그러면 위의 예제는 다음과 같이 될 것이다. 

 

132 392 311 351 231 -> 0 1 2 3 4

392 351 132 311 231 -> 1 3 0 2 4

origin B의 경우

	0	1	2	3	4
idxA	4	3	2	1	0
idxB	4	2	0	3	1

다음과 같이 정렬될 것이고 앞에서 부터 자신이 원래 B에 있던 인덱스에 넣어준 후 query(처음, 정렬 전 자기가 B에 있던 인덱스)를 구해주면된다. 그러면 자신보다 큰 것이 앞에오는 경우를 알 수 있다. 

 

위의 로직이 가능한 이유는 내림차 순으로 정렬을 해준 후 인덱스드 트리의 구성을 시작하기 때문에 가장 큰 원소를 넣었을 때는 1행에서의 가장 마지막 원소를 넣었기 때문에 자신의 뒤에 존재하는 것들이 자신의 앞에 나올 수 없다.(자신의 뒤가 없기 때문이다..!) 

두번째로 큰 원소를 넣었을 때는 자신보다 뒤에 나오던 원소만 체크해주면 된다(이미 앞에 들어가 있다!) 자신보다 앞은 신경안써줘도 된다. (인덱스드 트리에 들어있지 않다) 

이런 방식으로 인덱스드 트리에서는 넣는 순서로 자신보다 작은 것이나 큰 것들을 제외해줄 수 있다. 물론 인덱스드 트리의 업데이트 query 시간복잡도는 O(logN)이므로 훨씬 더 빠른 속도로 말이다!

 

---

4. 코드

#pragma warning (disable : 4996)
#include <cstdio>
#include <unordered_map>
#include <algorithm>
#define MAX 524288
#define SIZE_TR 1048576
#define ll long long
using namespace std;
struct BB {
	int idxA;
	int idxB;
	bool inline const operator < (const BB& ref) const {
		if (this->idxA > ref.idxA) return true;
		else return false;
	}
};
unordered_map<int, int> mp;
int N;
int originA[500002]; //int 배열 
BB originB[500002];
ll Tree[SIZE_TR];
int offset = 1;
ll final;
void input() {
	//입력을 받아들인다. 
	scanf("%d", &N);
	for (int i = 0; i < N; i++) {
		scanf("%d", &originA[i]);
		mp[originA[i]] = i; //인덱스를 저장 
	}
	for (int i = 0; i < N; i++) {
		int tmp;
		scanf("%d", &tmp);
		originB[i].idxA = mp[tmp]; // A인덱스를 저장한다. 
		originB[i].idxB = i; // B인덱스를 저장한다.
	}
}
int query(int from, int to) {
	int ans = 0;
	from += offset; to += offset;
	while (from <= to) {
		if (from % 2) ans += Tree[from++];
		if ((to % 2) == 0) ans += Tree[to--];
		from /= 2; to /= 2;
	}
	return ans-1;
}
void update(int idx) {
	idx += offset;
	Tree[idx] = 1;
	idx /= 2;
	while (idx > 0) {
		Tree[idx] = Tree[idx * 2] + Tree[idx * 2 + 1];
		idx /= 2;
	}
	/*for (int i = 0; i <= 15; i++) {
		printf("%d ", Tree[i]);
	}
	printf("\n");*/
}
void sol() {
	for (int i = 0; i < N; i++) {
		update(originB[i].idxB);
		final += query(0, originB[i].idxB);
	}
}
int main() {
	input();
	sort(originB, originB + N); // 오름차순으로 정렬 
	while (offset < N) offset *= 2; // offset을 만든다. 
	sol();
	printf("%lld\n", final);
	return 0;
}

마지막 함정은 정답이 int의 범위를 넘을 수 있다는 것이었다. 나는 이 점 때문에 계속 틀렸었다... 50만개와 50만개가 있을 경우 다 겹치는 상황에서 인트의 50만 ~ 1의 합이 int의 범위를 넘어가기 때문이다. 이 부분만 수정해주니 바로 맞았다.