BOJ 6198(G5)

https://www.acmicpc.net/problem/6198

은근히 발상이 어려워서 고민을 했다.

​

포인트는 한 건물에서 볼 수 있는 건물의 수를 따로따로 더하는 것이 아니라, 0~N-1까지 스캔하면서 해당 건물이 보이는 건물의 수를 더해준다고 생각하는 것이다. 예제를 가지고 순차적으로 사고해 보면, i번째 건물의 턴에서 자기보다 높이가 낮은 건물들은 치워주면 된다고 생각할 수 있다. 이때 순차적으로 스캔하기 때문에 최신(직전)의 건물들만 치워주면 높이 내림차순으로 자동 정렬되기에 추가적인 절차가 필요하지 않다. 후입선출이므로 스택을 사용한다.

 

import sys
input = sys.stdin.readline

N = int(input())
h = [int(input()) for _ in range(N)]

stack = []
res = 0
for i in range(N):
    while len(stack) > 0 and stack[-1] <= h[i]: stack.pop()
    stack.append(h[i])
    res += len(stack)-1

print(res)

 

 

그러나 처음에는 이 발상을 하지 못했다. 그래서 (자기 이상의 높이를 가진 건물 때문에 제약이 생기므로) 높이순으로 정렬한 다음, 높은 건물부터 하나씩 리스트에 넣어주면서 자기보다 크면서 최소 index를 가진 건물과의 거리를 구하면 된다고 생각했다.

문제는 구현인데, 평범한 리스트를 사용하면 '자기보다 큰 최소 index'를 찾는데 최대 $O(N)$이 소요되어 $O(N^2)$이 된다. index만 뽑아내서 이분 탐색을 하는 방법도 생각해 보았으나, 결국 자기의 index를 리스트에 삽입하는 데 $O(N)$이 걸린다.

그래서 update와 search가 모두 $log N$인 세그먼트 트리를 이용했다. tree의 값으로 index를 넣어준 다음, (x+1, N-1)구간 내 최소 index를 뽑아낼 수 있도록 했다. 이렇게 하면 시간복잡도는 $O(N log N)$이 되어 제약조건 내 풀이가 가능하다.

 

import sys
input = sys.stdin.readline


class SegTree:
    def __init__(self, size):
        self.N = size
        self.tree = [N]*(4*size)
    
    def update(self, node, l, r, target):
        if l==r:
            self.tree[node] = target
            return

        m = (l+r)//2
        if target <= m: 
            self.update(2*node, l, m, target)
        else: 
            self.update(2*node+1, m+1, r, target)
        
        self.tree[node] = min(self.tree[2*node], self.tree[2*node+1])
    
    def search(self, node, l, r, tl, tr):
        if tr < l or r < tl: return self.N
        if tl <= l and r <= tr: return self.tree[node]

        m = (l+r)//2
        return min(
            self.search(2*node, l, m, tl, tr),
            self.search(2*node+1, m+1, r, tl, tr)
        )
    

N = int(input())
segtree = SegTree(N)
h = [(int(input()), i) for i in range(N)]
h.sort(key=lambda x:x[0])
res = 0

for _ in range(N):
    height, idx = h.pop()
    min_idx = segtree.search(1, 0, N-1, idx+1, N-1)
    res += (min_idx - idx) - 1
    segtree.update(1, 0, N-1, idx)

print(res)

 

순차적 사고를 활용하는 아이디어가 구현에 있어 훨씬 간단하다는 교훈을 느낀다.