문제 설명

모든 건물에 대해, 다음 2가지를 구하라.
1. 좌측 & 우측에서 자기보다 크면서 다른 건물에 가리지 않은 건물의 개수
2. 그 건물 중에 가장 가까운 건물의 idx

"배열 원소 훑기 유형" 스택 문제가 꽤 많이 보이길래, 하나 정리해봤다.
이 문제는 해당 유형 중에서 조금 귀찮은 편이다.

아이디어

Naive $O(N^2)$ 아이디어

우선 Naive 한 완전 탐색 풀이부터 생각해보자.
모든 건물을 보면서, 각 건물마다 좌측 & 우측에 있는 조건에 맞는 건물을 전부 살펴보는 것.

전형적인 2중 반복문이므로 $O(N^2)$이 나올 것이다.

스택으로 $O(N)$으로 최적화

위 풀이를 스택을 이용해 $O(N)$으로 최적화 할 수 있다.
우선 건물을 왼쪽부터 보되, 각 건물마다 좌측에 있는 건물만 살펴보도록 하자. (우측은 따로 처리)

$i$번 건물 좌측에 있는 건물을 "전부" 살펴보는 대신, 왼쪽부터 훑으며 스택에 건물들을 넣어두고,
필요 없어진 저층 건물을 제외하고 필요한 건물만 스택에 남긴 채 확인하면 된다.

진행 과정

예제 입력 1 을 그대로 가져왔다.

8
3 7 1 6 3 5 1 7

  2           8
  2   4       8
  2   4   6   8
  2   4   6   8
1 2   4 5 6   8
1 2   4 5 6   8
1 2 3 4 5 6 7 8

먼저 $i=1$번 건물은 왼쪽에 아무것도 없으므로, 그냥 스택에 push(i=1)한다.

$i=2$번 건물의 왼쪽에는, $i=1$번 건물이 있다. (i.e. $i=1$이 스택에 들어있다.)
그런데 이 경우 $i=2$번 건물보다 $i=1$번 건물이 낮기 때문에, 이 건물은 조건에 맞지 않는다.

이번에 확인한, 낮은 $i=1$번 건물은 $i=2$번 건물에 가려서 $i\ge3$인 건물들한테 영원히 보일 일이 없다.
그러므로, 스택에서 $i=1$번 건물을 pop() 해버려도 문제가 없다.

그러면 스택이 비게 되므로, $i=2$ 좌측에 볼 수 있는 건물이 없음을 확인할 수 있다.
작업이 끝났으니 스택에 push(i=2)도 한다.

$i=3$번 건물은, 스택에 있는 $i=2$번 건물만 확인하면 된다.
이번엔 $i=2$번이 $i=3$번보다 높으므로, 스택에서 건물을 빼지 않는다.
이 때, 현재 스택에 들어있는 건물의 개수만큼, 좌측으로 볼 수 있는 건물이 존재함을 알 수 있다.
$i=3$번도 끝났으므로 push(i=3).

$i=4$번. 스택 위에는 $i=3$이, 아래에는 $i=2$가 들어있다.
스택 위에서부터 (i.e. 바로 왼쪽 건물부터) 보면, $i=3$이 $i=4$보다 낮기 때문에 스택에서 $i=3$ 을 제거한다.

아직 안 끝났다. 스택이 비거나, 자기보다 높은 건물이 나올 때까지 스택을 반복해서 본다.
$i=2$의 경우, $i=4$보다 높기 때문에 반복을 그만두고 스택 제거 작업을 멈춘다.
스택에는 $i=2$만이 들어있고, 이것만이 $i=4$에서 좌측으로 볼 수 있는 건물이다.
$i=4$도 끝. push(i=4).

마지막으로 $i=5$까지만 보자. 스택에는 위에 $i=4$, 아래에 $i=2$가 있다.
$i=4$가 $i=5$보다 높기 때문에, 스택 제거 작업을 멈춘다.
다시 말해, $i=5$에서는, 스택에 든 $i=4$와 $i=2$가 전부 보인다는 말이다.
이 때, 스택의 꼭대기에 있는 $i=4$가 가장 가까운 거리에 있는 보이는 건물이다.

핵심 논리

요컨데, 조건을 만족시키지 않는 $i$번보다 낮은 건물은
그 즉시 $i$번 건물에 가려서 더 이상은 볼 필요가 없어진다는 점을 이용하면,
다신 안 보도록 스택에서 빼버리는 풀이가 가능하다.

우측으로 보이는 건물도 마찬가지 방법으로 찾을 수 있다.
좌측 풀이를 거울에 비춘듯 하면 된다.
우측부터 좌측으로 건물을 훑어보며, 우측으로 보이는 것들을 스택에 넣으며 판정.

핵심 코드

N = int(input().rstrip())
A = list(map(int, input().split()))
st = []  # id

ans_cnt = [0] * N
ans_min_dist_id = [INF] * N

# i번 건물의 좌측 관찰
for i in range(N):
    # i번 건물에 가린 낮은 건물들은, 뒷 건물에서도 영원히 못본다.
    while len(st) != 0 and A[st[-1]] <= A[i]:
        st.pop()
    if len(st) != 0:
        # 스택에 남은 건물의 개수가, i번 왼쪽에서 볼 수 있는 건물의 수
        ans_cnt[i] = len(st)
        # 스택의 꼭대기 건물이 왼쪽에서 가장 가까이 있는 볼 수 있는 건물
        ans_min_dist_id[i] = st[-1]
    st.append(i)

시간 복잡도

이 풀이의 시간 복잡도는 $O(N)$이다.

위 코드를 보면, 일단 for 문은 N번 수행됨은 자명하다.
문제는 내부 while 문인데, 각 $i$마다 반복 횟수는 입력에 따라 다르다.

하지만, 프로그램 전체를 통틀어 while 문 내부 st.pop()이 실행되는 횟수가
최대 N번이라는 건 알 수 있다.

왜냐하면 st.append(i)는 총 N번 호출되니, st엔 총 N개의 정수가 들어가는데,
그러면 당연히 N번 넘게 스택에서 꺼내는 건 불가능하기 때문.

헤맨 점

저장할 정보가 뭔지, 필요없는 정보가 뭔지 헷갈려서 오래 걸렸다.

1. ans 배열에 거리는 따로 저장할 필요가 없다. ans[i]에 id 저장 후 abs(id-i)하면 거리다.
2. 스택에 높이를 따로 저장할 필요가 없다. id 저장 후 A[id]로 구하면 그만이다.

입력, 필요한 중간 변수, 출력이 뭔지는 바로바로 보여야 하는데, 아직도 미숙한가 싶다.

정답 코드 : $O(N)$