[PS] 오아시스 재결합 / C++ (백준 3015번)

오아시스 재결합 (백준 3015번)

 

보자마자 히스토그램 문제가 생각나서 스택으로 풀어보려다가 너무 뻔한 것 같아서 정렬, 이분 탐색으로 풀어봤다.

 

시간 초과 코드($O(n^2)$)

• vector의 insert 연산이 $O(n)$이라는 것을 까먹고 짰다가 시간 초과를 받은 코드
• 키 순서대로 정렬한 다음 벡터에 이진 탐색으로 키와 index를 삽입하면서 좌, 우에 나보다 더 큰 수가 있으면 dap++ 해준다.
  • 똑같은 수가 연속적으로 삽입될 경우, 개수를 세어 "$kC2$ + $k$(왼쪽에 나보다 큰 수가 존재할 경우) + $k$(오른쪽에 나보다 큰 수가 존재할 경우)"를 더해주었다.
• 제미나이를 사용해 이를 세그먼트 트리로 구현하여 $O(nlog^2n)$으로 줄여서 제출해 봤지만 시간 초과가 나왔다.
  • n이 최대 500,000이어서 연산 횟수가 최대 $500,000 \times (\log_2 500,000)^2 \approx 500,000 \times 19^2 \approx 1.8$억 번이면 1초가 넘기 때문

#include<iostream>
#include<algorithm>
#include<vector>
using namespace std;
typedef long long ll;

struct data {
	int tall, index;
};

pair<int, int> pa[500'001];
vector<struct data> vec;

int comp(pair<int, int> pa1, pair<int, int> pa2) {
	if (pa1.first == pa2.first)return pa1.second > pa2.second;
	return pa1.first > pa2.first;
}

int main()
{
	int n;
	ll dap = 0;
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) {
		scanf("%d", &pa[i].first);
		pa[i].second = i;
	}
	sort(pa + 1, pa + n + 1, comp);
	for (int i = 1; i <= n; i++) {
		int st = 0, en = vec.size();
		if (en < 1)vec.push_back({ pa[i].first, pa[i].second});
		else {
			int mid = (st + en) / 2;
			while (st < en) {
				if (vec[mid].index < pa[i].second)st = mid + 1;
				else en = mid;
				mid = (st + en) / 2;
			}
			int max_i = mid, min_i = mid;
			ll k = 1;
			vec.insert(vec.begin() + mid, { pa[i].first, pa[i].second});
			while (i + 1 <= n && pa[i].first == pa[i + 1].first) {
				if (mid - 1 >= 0 && vec[mid - 1].index > pa[i + 1].second)break;
				i++; k++;
				vec.insert(vec.begin() + mid, { pa[i].first, pa[i].second});
				max_i++;
			}
			int size = vec.size();
			if (k == 1) {
				if (mid + 1 < size)dap++;
				if (mid - 1 >= 0)dap++;
			}
			else {
				if (max_i + 1 < size)dap += k;
				if (min_i - 1 >= 0)dap += k;
				dap += (k * (k - 1)) / 2;
			}
		}
	}
	printf("%lld\n", dap);
	return 0;
}

 

풀이 과정

• 결국 정해는 스택을 사용하는게 맞았다. 스택에 {키, 연속적으로 키가 같은 사람 수}를 다음 세 가지 상황별로 삽입해주면 된다.
  1. st.top() > next
     • st.top()에 있는 사람과 한 쌍이 되기 때문에 dap++ 해준다.
     • 스택에 {next, 1}를 삽입하고 다음 사람으로 넘어간다.
  2. st.top() == next
     • 같은 키인 사람이 st.top().second만큼 있다면 dap+=st.top().second 해준다.
       • ex) 키가 9인 사람이 연속 3명이 나왔었다면 3명과 모두 한 쌍을 이룰 수 있기 때문에 +3 해줌
     • 스택 size가 2 이상이라면 이전에 키가 더 큰 사람이 있었다는 것이므로 dap++ 해준다.
       • ex) {{13, 2}, {9, 3}} 에서 9가 한 번 더 나왔다면 맨 마지막 13과 한 쌍이 되며, 그전에 나온 13과는 한 쌍이 되지 못하기 때문에 1만 더해줌
     • st.top().second++ 하여 키가 똑같은 사람이 연속으로 한 명 더 늘었다는 것을 표현해 준다.
  3. st.top() < next
     • 더 큰 값인 next로 인해 이 뒤로 쌍을 맺지 못하므로 next와 쌍을 이루는 것 만 더해준다.
     • st.top().second()를 더해야 하며 이땐 break나 push 하지 않고 pop만 해준다.
       • ex) {{13, 2}, {9, 4}}에서 10이 나왔을 경우 키가 9인 사람 4명과 10은 모두 짝을 맺을 수 있으므로 +4 해주고, pop 한 다음 13과 비교하게끔 break 하지 않음
• 결국 스택을 내림차순으로 유지하는 모노톤 스택을 구현하는 내용이었다.

#include<iostream>
#include<algorithm>
#include<stack>
using namespace std;
typedef long long ll;

int main()
{
	int n;
	ll dap = 0;
	stack<pair<int, int> > st;
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) {
		int next;
		scanf("%d", &next);
		while (1) {
			if (st.empty()) {
				st.push(make_pair(next, 1));
				break;
			}
			else {
				if (st.top().first > next) {
					st.push(make_pair(next, 1));
					dap++;
					break;
				}
				else if (st.top().first == next) {
					dap += st.top().second;
					if (st.size() > 1)dap++;
					st.top().second++;
					break;
				}
				else {
					dap += st.top().second;
					st.pop();
				}
			}
		}
	}
	printf("%lld\n", dap);
	return 0;
}