[C++][백준] 크게 만들기(2812) - 세그먼트 트리, 스택
문제
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;
}