[정올/JUNGOL] 3234번 화살표(고) (C/C++)

https://jungol.co.kr/problem/3234

* 사용언어 : C언어, C++

 

1. 문제

N개의 색깔 중 하나의 색을 가지는 N개의 점이 직선 위에 주어짐

모든 점에 대해서 자신과 같은 색인 점 중 가장 가까운 곳까지의 길이의 합을 출력

같은 색인 점이 없으면 길이는 0으로 처리하고, 2개 이상이면 아무거나 선택

 

2. 풀이

모든 점으로부터 자신과 같은 색인 점을 찾고, 그 안에서 가장 가까운 점을 찾아야 합니다. 이 작업을 각각의 점에 대해서 매번 전체 순회하면서 찾으면 시간 복잡도가 O(N^2)이 되므로 TLE가 날 것 같습니다. (제출 해보진 않았습니다)

 

대신 전체를 아래 우선순위에 따라 정렬해두면 좌우에 있는 점만 확인하면 됩니다.1) 색깔 오름차순 (내림차순도 가능)2) 좌표 오름차순 (내림차순도 가능)

 

좌우에 있는 점이 유효하지 않은 경우(좌표 0 미만 or 좌표 N 이상 or 색이 다른 경우)

INF(1<<30)를 return하게 했고 좌우 모두 INF인 경우 0으로 처리하게 하면 됩니다.

 

그 외 설명은 아래 코드로 대체하겠습니다.

 

3. 코드

#ifndef _CRT_SECURE_NO_WARNINGS
#define _CRT_SECURE_NO_WARNINGS
#endif

#include <stdio.h>

#define ABS(X) (((X) > 0) ? (X) : (-(X)))
#define MIN(A, B) (((A) < (B)) ? (A) : (B))
#define INF (1 << 30)

const int LM = 100000;
int N;

struct _p {
	int pos, color;

	bool operator< (const _p &other) const {
		if (color == other.color) return pos < other.pos;
		return color < other.color;
	}
} P[LM], *arr[LM], *tmp[LM];

void mergeSort(int s, int e) {
	if (s >= e) return;
	int m = (s + e) / 2;

	mergeSort(s, m);
	mergeSort(m + 1, e);

	int i = s, j = m + 1, k = s;
	while (i <= m && j <= e) {
		if (*arr[i] < *arr[j]) tmp[k++] = arr[i++];
		else tmp[k++] = arr[j++];
	}

	while (i <= m) tmp[k++] = arr[i++];
	while (j <= e) tmp[k++] = arr[j++];

	for (i = s; i <= e; ++i) arr[i] = tmp[i];
}

int getMinDist(int i, int dir) {
	int j = i + dir;
	if (j < 0 || j >= N) return INF;
	if (arr[i]->color != arr[j]->color) return INF;
	return ABS(arr[i]->pos - arr[j]->pos);
}

int main() {
	scanf("%d ", &N);
	for (int i = 0; i < N; ++i) {
		scanf("%d %d", &P[i].pos, &P[i].color);
		arr[i] = &P[i];
	}

	mergeSort(0, N - 1);

	long long res = 0;
	for (int i = 0; i < N; ++i) {
		int left = getMinDist(i, -1);
		int right = getMinDist(i, 1);

		if (left == INF && right == INF) continue;
		res += MIN(left, right);
	}

	printf("%lld\n", res);
	return 0;
}