[C++ 백준 풀이] 1849번 : 순열 / 1777번 : 순열복원 (세그먼트 트리)

이 포스트에서는 프로그래밍 문제 사이트 백준 Online Judge의 1849번 : '순열', 1777번 : '순열복원' 문제에 대한 풀이 코드와 해설을 다루고 있습니다.

 

문제의 난이도는 Solved.ac 기준 Platinum V에 해당하며, 문제를 풀이하기 위해 세그먼트 트리의 응용 개념을 다룰 것입니다.

 

또한 1849번 : 순열 문제와 1777번 : 순열 복원 문제는 내용 상 연결이 되는 문제이고 같은 방식으로 풀리기 때문에 묶어서 다루도록 하겠습니다.

 

 

1849번 : 순열

 

1~N 사이의 숫자가 한 번씩 등장하는 수열에 대해, 숫자 i 앞에 위치한 수들 중 i보다 큰 수들의 개수 A[i]에 대한 순열이 주어지면 원래 수열을 찾는 문제입니다.

간단히 원리를 생각해보면, 1부터 위치를 시키되 A[i]만큼 앞에 빈 칸을 만들고 수를 위치시키면 됩니다.

 

위의 예제로 예시를 들면, 1은 앞에 5개의 큰 수가 있으므로 _ _ _ _ 1으로 먼저 배치를 시킵니다.

그 다음 2는 앞에 큰 수가 없으므로 가장 앞에 와야합니다. (예외가 될 수 있는 것은 1인데 2보다 작은 수들을 먼저 배치시켰기 때문에 고려할 필요가 없어짐)

그러면 2 _ _ _ 1과 같이 배치가 됩니다.

그 다음 3은 앞에 큰 수가 1개라고 했으므로 1개의 빈 칸을 두고 2 _ 3 1과 같이 배치하면 됩니다.

이러한 방식으로 배치를 하면 2 7 3 5 4 1 8 6의 순열을 만들 수 있습니다.

 

 

 

#include <cstdio>
#include <vector>
using namespace std;
 
vector<int> Tree, Arr;
 
void initTree(int Node, int Begin, int End) {
    if(Begin == End) {
        Tree[Node] = 1;
        return;
    }
    int Mid = (Begin + End)/2;
    initTree(Node*2, Begin, Mid);
    initTree(Node*2+1, Mid+1, End);
    Tree[Node] = Tree[Node*2] + Tree[Node*2+1];
}
 
int findIndex(int Node, int Begin, int End, int Count) {
    if(Begin == End) return Begin;
    int Mid = (Begin + End)/2;
    if(Count < Tree[Node*2]) return findIndex(Node*2, Begin, Mid, Count);
    else return findIndex(Node*2+1, Mid+1, End, Count-Tree[Node*2]);
}
 
void updateTree(int Node, int Begin, int End, int Index) {
    if(Index < Begin || Index > End) return;
    if(Begin == End) {
        Tree[Node] = 0;
        return;
    }
    int Mid = (Begin + End)/2;
    updateTree(Node*2, Begin, Mid, Index);
    updateTree(Node*2+1, Mid+1, End, Index);
    Tree[Node]--;
}
 
int main() {
    int N, Count, Index;
    scanf("%d", &N);
    Tree.resize(N*4+1), Arr.resize(N+1);
    initTree(1, 1, N);
    for(int i=1; i<=N; i++) {
        scanf("%d", &Count);
        Index = findIndex(1, 1, N, Count);
        Arr[Index] = i;
        updateTree(1, 1, N, Index);
    }
    for(int i=1; i<=N; i++) printf("%d\n", Arr[i]);
}
 

이러한 알고리즘을 O(N^2)보다 작게 구현하기 위해서는 세그먼트 트리를 사용하여, 구간의 빈 칸의 수를 저장하도록 해주면 됩니다.

구현은 모든 리프 트리에 1을 저장해주고 (초기에는 모두 빈 칸이므로) 값을 하나 채울 때마다 그 리프 노드를 0으로 만들어서 갱신해주면 됩니다.

핵심은 findIndex 함수의 조건문이라고 생각하는데, Count와 왼쪽 트리의 빈칸의 합인 Tree[Node*2]를 비교해서 더 작은 경우는 왼쪽 트리로, 그 외에는 오른쪽 트리로 이동하면서 Tree[Node*2]를 빼고 Count 해줘야합니다.

Count와 왼쪽 트리의 합이 같을 때에는 오른쪽 서브트리로 이동해야 "앞의" 빈칸이 Count개인 Index를 반환해줄 수 있습니다.

 

 

 

풀이를 제출해보면 굉장히 짧은 시간에 문제를 해결할 수 있음을 확인할 수 있습니다.

풀이 시간이 넉넉하기 때문에 세그먼트 트리를 이용하지 않고 vector만로도 아슬아슬한 시간에 해결할 수도 있다고 합니다.

 

 

1777번 : 순열 복원

 

위의 문제와 거의 비슷한 순열 복원 문제입니다.

다른 점은 앞이 아닌 뒤에 나오면서 더 작은 수의 개수 A[i]가 주어집니다.

 

비슷한 원리로 순서만 반대로 해서 위의 예제로 생각해봅시다.

처음에 가장 큰 숫자 8을 뒤에서부터 빈칸이 0개 있도록 배치하면 맨 뒤에 배치됩니다.

그 다음 숫자 7을 뒤에서부터 빈칸이 2개 존재하도록 배치하면 7 _ _ 8이 됩니다.

그 다음 숫자 6을 뒤에서부터 빈칸이 1개 존재하도록 배치하면 7 6 _ 8이 됩니다.

같은 방법으로 순열을 모두 채우면 2 4 5 1 7 6 3 8이 됩니다.

 

이제 이 과정을 똑같이 코드로 구현하기만 하면 됩니다.

 

 

#include <cstdio>
#include <vector>
using namespace std;
 
vector<int> Tree, InvSeq, Arr;
 
void initTree(int Node, int Begin, int End) {
    if(Begin == End) {
        Tree[Node] = 1;
        return;
    }
    int Mid = (Begin + End)/2;
    initTree(Node*2, Begin, Mid);
    initTree(Node*2+1, Mid+1, End);
    Tree[Node] = Tree[Node*2] + Tree[Node*2+1];
}
 
int findIndex(int Node, int Begin, int End, int Count) {
    if(Begin == End) return Begin;
    int Mid = (Begin + End)/2;
    if(Count < Tree[Node*2]) return findIndex(Node*2, Begin, Mid, Count);
    else return findIndex(Node*2+1, Mid+1, End, Count-Tree[Node*2]);
}
 
void updateTree(int Node, int Begin, int End, int Index) {
    if(Index < Begin || Index > End) return;
    if(Begin == End) {
        Tree[Node] = 0;
        return;
    }
    int Mid = (Begin + End)/2;
    updateTree(Node*2, Begin, Mid, Index);
    updateTree(Node*2+1, Mid+1, End, Index);
    Tree[Node]--;
}
 
int main() {
    int N, Count, Index;
    scanf("%d", &N);
    Tree.resize(N*4+1), InvSeq.resize(N+1), Arr.resize(N+1);
    initTree(1, 1, N);
    for(int i=1; i<=N; i++) scanf("%d", &InvSeq[i]);
    for(int i=N; i>=1; i--) {
        Index = findIndex(1, 1, N, InvSeq[i]);
        Arr[Index] = i;
        updateTree(1, 1, N, Index);
    }
    for(int i=N; i>=1; i--) printf("%d ", Arr[i]);
}
 

 

코드를 살펴보면 거의 비슷하고, 메인 함수의 루프 방향만 반대임을 알 수 있습니다.

우선 큰 수부터 탐색을 수행해야 하기에, InvSeq 배열을 따로 선언하여 뒤에서부터 처리할 수 있도록 하였습니다.

그리고 출력부 역시 큰 수부터 역으로 수행했기에, Arr 배열 역시 역순으로 출력해주어야 합니다.

나머지 과정은 거의 비슷합니다.

 

 

 

풀이 코드를 제출해보면 짧은 시간에 모든 테스트케이스를 통과하여 정답 처리를 받을 수 있음을 확인할 수 있습니다.