카테고리 없음

[C++][백준] 크게 만들기(2812) - 세그먼트 트리, 스택

Belluga 2021. 5. 13. 15:36

문제

N자리 숫자가 주어졌을 때, 여기서 숫자 K개를 지워서 얻을 수 있는 가장 큰 수를 구하는 프로그램을 작성하시오.

 

출력

입력으로 주어진 숫자에서 K개를 지웠을 때 얻을 수 있는 가장 큰 수를 출력한다.

 

풀이

이 문제는 세그먼트 트리, 스택(그리디)이라는 두가지 방식을 이용하여 풀어보았습니다.

 

먼저 스택을 이용한 풀이 방식입니다.

199999보다 911111가 큰 수 입니다. 가장 큰 수를 얻기 위해서는 앞 자리수에 최대한 큰 수를 배치해야 합니다.

 

따라서 스택에 순차적으로 수를 집어넣다가

현재 스택 top보다 큰 수가 존재하는 경우 delete_cnt가 허용하는 내로 숫자를 지워가며 최대한 앞자리에 큰 수가 올 수 있도록 만들어줍니다.

#include <stack>
#include <iostream>
#include <string>

using namespace std;

const int MAX = 500010;

int N, K;
string s;
int arr[MAX];

int main() {
	cin >> N >> K >> s;

	for (int i = 0; i < N; i++) {
		arr[i] = s[i] - '0';
	}

	stack<int> myStack;
	int delete_cnt = 0;

	for (int i = 0; i < N; i++) {
		while (true) {
			if (myStack.empty()) {
				myStack.push(arr[i]);
				break;
			}
			if (myStack.top() < arr[i] && delete_cnt < K) {
				delete_cnt++;
				myStack.pop();
			}
			else {
				myStack.push(arr[i]);
				break;
			}
		}
	}

	stack<int> myStack2;
	while (!myStack.empty()) {
		myStack2.push(myStack.top());
		myStack.pop();
	}

	for (int i = 0; i < N - K; i++) {
		cout << myStack2.top();
		myStack2.pop();
	}

	return 0;
}

마지막으로 입력으로 주어진 숫자에서 K개를 무조건 지워야 하기 때문에 N-K개까지만 출력하였습니다. 

 

참조 테스트 케이스

5 2
54321 // output : 543

 

다음으로 세그먼트 트리를 이용한 풀이 방식입니다.

 

참조 테스트 케이스

9 5
192456017

 

아홉자리 수에서 5개의 수를 제거하여 가장 큰 수를 만들어야 합니다.

즉 가장 큰 네자리 수를 만들어야 합니다.

가장 큰 네자리 수를 만들기 위해 인덱스 1~6 범위에서 최댓값을 찾습니다.

네자리 수를 만들어야 하기 때문에 마지막 숫자 3개는 남겨두었습니다.

해당 범위에서 최댓값은 9이고 이는 정답의 첫번째 숫자가 됩니다.

 

※ 참고로 가장 큰 수가 범위 내에 여러개가 있는 경우 인덱스가 더 작은 값을 최댓값으로 선택합니다. 

(그래야 다음 자리에서도 큰 수를 고를 수 있습니다)

 

 

숫자 9 다음의 인덱스 3~7 범위에서 최댓값을 찾습니다.

해당 범위에서 최댓값은 6이고 이는 정답의 두번째 숫자가 됩니다.

 

 

숫자 6 다음의 인덱스 7~8 범위에서 최댓값을 찾습니다.

해당 범위에서 최댓값은 1이고 이는 정답의 세번째 숫자가 됩니다.

 

 

숫자 1 다음의 인덱스 9 에서 최댓값을 찾습니다.

해당 범위에서 최댓값은 7이고 이는 정답의 마지막 숫자가 됩니다.

 

특정 구간에 대해 최댓값을 찾을 때는 세그먼트 트리를 이용할 수 있습니다.

 

초기 배열을 기반으로 세그먼트 트리를 만드는 init 함수l, r 구간에 대해 최댓값을 찾는 query 함수입니다.tree[node].first는 value, tree[node].second는 index를 의미합니다.

pair<int, int> init(int node, int l, int r) {
	// leaf
	if (l == r) {
		tree[node].first = arr[l];
		tree[node].second = l;
		return tree[node];
	}

	int mid = (l + r) / 2;
	
	return tree[node] = max(init(node * 2, l, mid), init(node * 2 + 1, mid + 1, r));
}

pair<int, int> query(int l, int r, int node, int nl, int nr) {
	if (l > nr || r < nl) return make_pair(-1, MAX + 10);

	if (nl >= l && nr <= r) return tree[node];

	int mid = (nl + nr) / 2;
	return max(query(l, r, node * 2, nl, mid), query(l, r, node * 2 + 1, mid + 1, nr));
}

 

if (l > nr || r < nl) return make_pair(-1, MAX + 10);

이때 query 함수의 위 if문 같은 경우 습관처럼 0을 return 하면 안됩니다.

구간 내 최댓값이 0이 될 수도 있는데 우리는 인덱스가 더 작은 값을 최댓값으로 선택합니다.

따라서 값을 -1, 인덱스를 MAX + 10으로 return 하도록 했습니다.

 

전체 코드는 아래와 같습니다.

#include <iostream>
#include <string>
using namespace std;

const int MAX = 500010;
int arr[MAX];
pair<int, int> tree[MAX * 4];

int N, K, remain;
string string_input;

pair<int, int> max(pair<int, int> a, pair<int, int> b) {
	if (a.first == b.first) {
		if (a.second < b.second) {
			return a;
		}
		else {
			return b;
		}
	}
	else {
		if (a.first < b.first) {
			return b;
		}
		else {
			return a;
		}
	}
}

pair<int, int> init(int node, int l, int r) {
	// leaf
	if (l == r) {
		tree[node].first = arr[l];
		tree[node].second = l;
		return tree[node];
	}

	int mid = (l + r) / 2;
	
	return tree[node] = max(init(node * 2, l, mid), init(node * 2 + 1, mid + 1, r));
}

pair<int, int> query(int l, int r, int node, int nl, int nr) {
	if (l > nr || r < nl) return make_pair(-1, MAX + 10);

	if (nl >= l && nr <= r) return tree[node];

	int mid = (nl + nr) / 2;
	return max(query(l, r, node * 2, nl, mid), query(l, r, node * 2 + 1, mid + 1, nr));
}

int main() {
	cin >> N >> K;
	remain = N - K;

	cin >> string_input;
	for (int i = 0; i < N; i++) {
		arr[i + 1] = string_input[i] - '0';
	}	
	init(1, 1, N);

	int return_idx = 0;
	for (int i = 0; i < remain; i++) {
		pair<int, int> return_value = query(return_idx + 1, N - remain + 1 + i, 1, 1, N);

		cout << return_value.first;
		return_idx = return_value.second;
	}

	return 0;
}