[백준] 22866 : 탑 보기 - Python

문제

N개의 건물이 일렬로 세워져 있습니다. 각 건물에서 볼 수 있는 다른 건물의 수와 가장 가까운 건물의 번호를 구하는 문제입니다. 단, 자신보다 높은 건물만 볼 수 있으며, 가장 가까운 건물이 여러 개일 경우 건물 번호가 작은 것을 선택해야 합니다.

제약 조건:

• 1 ≤ N ≤ 100,000
• 1 ≤ 건물의 높이 ≤ 1,000,000,000

내 코드

import sys
sys.stdin = open("input.txt", "rt")

# 탑 보기
# 건물 n개, 각 건물 양 옆에 보이는 건물 수 구하는 문제
# 본인 높이보다 큰 건물만 볼 수 있다.
# n <= 100,000 -> N^2 안됨
# 각 건물에서 볼 수 있는 건물 개수 구하기,
# 만약 볼 수 있는 건물 여러개 -> 거리가 가장 가까운 건물의 번호 중 작은 번호


n = int(input())
heights = list(map(int, input().split())) # 건물들의 높이
# -> 보자마자 스택 떠올림
cnt = [0] * n
stack = []
INF = int(1e9)
nearest = [INF] * n # 가장 가까운 건물의 번호 저장
# 먼저 왼쪽에서 오른쪽으로 진행.
for i in range(n):
    now_height = heights[i] # 현재 건물 높이
    while stack and heights[stack[-1]] <= now_height: # 스택에서 건물의 인덱스 저장.
        stack.pop()
    cnt[i] += len(stack) # 본인보다 왼쪽에 있는 건물 중 높이 더 높은 건물 개수 저장
    if stack: # 본인보다 높은 건물이 있을 때만 갱신
        nearest[i] = stack[-1] # 현재 건물의 가장 가까운 인덱스 저장 (오른쪽에서 돌 때 비교 진행)
    stack.append(i)

# 오른쪽에서 왼쪽으로 진행
stack = []
for i in range(n-1, -1, -1):
    now_height = heights[i] # 현재 건물 높이
    while stack and heights[stack[-1]] <= now_height:
        stack.pop()
    cnt[i] += len(stack)
    # 이제 여기서 가장 가까운 건물의 인덱스 넘버 비교해야해.
    # 1. 거리가 가까운, 거리가 같다면 건물 번호가 더 작은
    # 현재 건물과 기존 저장된 가장 가까운 건물 거리와 현재 건물과 여기서의 가장 가까운 건물 거리 비교
    if stack:
        if nearest[i] == INF: # 현재 가장 가까운 게 없다면
            nearest[i] = stack[-1]
        elif i - nearest[i] > stack[-1] - i: # 지금의 거리가 더 작은 경우에 교체
            nearest[i] = stack[-1]
        elif i - nearest[i] == stack[-1] - i: # 거리가 같은 경우
            nearest[i] = min(nearest[i], stack[-1]) # 더 작은 번호로 저장
    stack.append(i)

for i in range(n):
    if cnt[i] == 0:
        print(0)
    else:
        print(f"{cnt[i]} {nearest[i] + 1}")

코멘트

1. 처음 생각과 틀린 이유:
   처음에는 각 건물에 대해 양쪽으로 순회하며 볼 수 있는 건물을 찾으려 했습니다. 하지만 이 방법은 O(N^2)의 시간 복잡도를 가져 시간 초과가 발생했습니다.

2. 스택을 떠올린 이유:
   효율적인 해결을 위해 스택을 사용하면 한 번의 순회로 각 건물에서 볼 수 있는 다른 건물들을 파악할 수 있다는 것을 깨달았습니다.

3. 문제 접근:

   • 왼쪽에서 오른쪽으로, 오른쪽에서 왼쪽으로 두 번 순회합니다.
   • 스택에는 현재 건물보다 높은 건물의 인덱스만 유지합니다.
   • 각 방향에서 스택의 크기가 볼 수 있는 건물의 수가 됩니다.
   • 가장 가까운 건물은 스택의 top에 있는 건물입니다.

4. 왜 스택을 써야 했는지:
   스택을 사용하면 각 건물에 대해 O(1)의 시간으로 볼 수 있는 건물의 수와 가장 가까운 건물을 찾을 수 있습니다. 이는 전체적으로 O(N) 시간 복잡도를 달성하게 해줍니다.

5. 얻은 인사이트:

   • 초기화의 중요성: nearest 리스트를 0이 아닌 INF로 초기화해야 정확한 비교가 가능했습니다. 처음에 nearest에 전부 0으로 해놨더니, 이후에 거리 계산에서 왜 틀린지 몰랐던..
   • 양방향 순회의 효율성: 한 번의 왼쪽 순회와 한 번의 오른쪽 순회로 모든 정보를 얻을 수 있었습니다.
     현재보다 높은 건물 유지 -> + N^2의 알고리즘 쓸 수 없다. -> 스택을 고려해야 하는 인사이트 얻음