[백준/1849] 순열 (Python)

Problem : https://www.acmicpc.net/problem/1849

1849번: 순열

 

Difficulty : Platinum 5

 

Status : Solved

 

Time : 00:31:05

 

---

 

문제 설명

 

 

1부터 N까지의 수들이 한 번씩 쓰인 수열이 있다. 그 수열에서 i 앞에 있는 수 들 중, i보다 큰 수들의 개수를 A[i]라고 정의하자. A[i]가 주어져 있을 때, 원래 수열을 구하는 프로그램을 작성하여라. 

 

 

입력 및 출력

 

입력

 

첫째 줄에 수열 원소의 개수 (1 ≤ N ≤100,000)이 주어진다. 그리고 두 번째 줄부터 N+1 번째 줄까지 차례로 A[i]가 주어진다.

 

출력

 

N개의 줄에 걸쳐서 수열을 출력한다. (i번째 줄에는 i번째 원소를 의미한다)

 

입력 예시

 

8
5
0
1
2
1
2
0
0

 

출력 예시

 

2
7
3
5
4
1
8
6

 

 

---

 

풀이

 

i번째 수에 대해 A[i]가 주어졌다면, 아직 배정되지 않은 자리들 중 A[i]+1번째에 i번째 수가 위치할 수 있다. 즉 아직 배정되지 않은 자리를 빠르게 구할 수 있어야 한다는 전제가 붙는다. 이러한 탐색에 O(logN)이 소요되는 이분 탐색, 혹은 세그먼트 트리를 사용해보자.

 

지금 이 문제에 주어진 시간제한은 0.5초. 기존 풀이코드는 재귀식으로 세그먼트 트리의 search와 update를 사용하였기에 필히 시간초과를 겪을 수밖에 없다. (python의 함수 호출은 상상 이상으로 느리다) 즉 재귀 형태로 정의된 함수들을 반복문 형태로 바꾸어 주어야 한다. 세그먼트 트리의 search가 구간합이 아닌 k+1번째 합을 구하는 경우이므로 이분 탐색 문제로 바꿀 수 있으며, update 시 단일 노드와 그 노드의 직접적인 조상 노드들에만 영향을 끼친다는 사실을 눈여겨 보자. start, end를 이분 탐색 문제처럼 초기화하고 조건에 맞을 때까지 update하기만 하면 된다.

 

풀이 코드

import sys
input = sys.stdin.readline

N = int(input())

class SegTree :
  def __init__(self) :
    self.tree = [0]*(4*N)
    self.result = [0]*N
    self.index = 1
    def _init(start, end, idx) :
      if start == end :
        self.tree[idx] = 1
        return
      mid = (start + end) // 2
      _init(start, mid, idx*2)
      _init(mid+1, end, idx*2+1)
      self.tree[idx] = self.tree[idx*2] + self.tree[idx*2+1]
    _init(0, N-1, 1)

  def search(self, rank) :
    start, end, idx = 0, N-1, 1
    while start < end :
      mid = (start + end) // 2
      if self.tree[idx*2] > rank :
        end = mid
        idx = idx*2
      else :
        start = mid + 1
        rank -= self.tree[idx*2]
        idx = idx*2+1
    return end

  def update(self, target) :
    start, end, idx = 0, N-1, 1
    self.tree[idx] -= 1
    while start < end :
      mid = (start + end) // 2
      if mid >= target :
        end = mid
        idx = idx*2
      else :
        start = mid + 1
        idx = idx*2+1
      self.tree[idx] -= 1
    self.result[target] = self.index
    self.index += 1

segtree = SegTree()
for _ in range(N) :
  num = int(input())
  target = segtree.search(num)
  segtree.update(target)
print(*segtree.result, sep = '\n')

풀이 완료!