Stack 알고리즘

Stack, Queue은 자료구조를 시작했을 때, 가장 처음 만나는 자료구조이다.

 

하지만.. 알고리즘에서 Stack 유형 알고리즘을 풀 때에는 번번히 어려움을 마주한다..😂

먼저 Stack 유형이라는 것을 파악하기에도 힘들고(DP / Greedy로 우선 접근), 파악한다고 하더라도 응용 문제에서는 자체 테스트케이스에서 틀리는 경우가 많다.

 

물론 지금까지 Stack 문제를 많이 풀지 않아서 그런 것 같기도 하지만..

무언가 정답 코드에 테스트 예제를 디버깅하면서 추적하면 이해가 잘 되지만, 그 전에는 확 와닿지 않는 느낌이랄까

 

그렇지만 Stack 문제를 풀 때마다 새롭고, 재밌다🤩

일반 구현에 비해 직관적(눈은 직관적이지만, 머리는 복잡한..)이고, 시간복잡도 측면에서도 대부분 O(N)으로 효율적이게 해결할 수 있다.

 

여러 Stack 유형의 문제들을 리뷰하면서 더 친숙해져보도록 하자!

 

 

---

1. 문제 예시

*BOJ - 17289. 오큰수 (골드 4)

 

[문제]

수열의 각 원소 Ai에 대해서 오큰수 NGE(오른쪽에 있으면서 Ai보다 큰 수 중에서, 가장 왼쪽에 있는 수) 리스트 값을 구한다.

 

[코드]

N = int(input())
lst = list(map(int, input().split()))
NGE = [-1] * N

idx = -1 # 0부터 시작하기 위함
indexStack = [] # 오큰수를 찾지 못하고 있는 값들의 인덱스 리스트
for num in lst:

    idx += 1
    while indexStack and lst[indexStack[-1]] < num:
        NGE[indexStack.pop()] = num

    indexStack.append(idx)

print(*NGE)

 

[풀이]

문제의 핵심은 idx 이전까지의 원소 중에서, 오큰수를 찾지 못하고 있는 값들의 인덱스를 stack에 관리한다. 그리고 stack의 top 원소의 숫자와, 현재 숫자의 크기를 비교하여 오큰수 조건을 만족한 경우, 남아 있는 stack 값들 또한 조건을 만족하는지 확인하며 오큰수를 찾아주도록 한다.

 

 

---

2. 문제 적용

1) BOJ - 17299. 오등큰수 (골드 3)

 

[문제]

수열의 각 원소 Ai에 대해서 오큰수 NGF(오른쪽에 있으면서 수열에서 등장한 횟수가 Ai보다 많은 수 중에서, 가장 왼쪽에 있는 수) 리스트 값을 구한다. 

 

[코드]

from collections import Counter

N = int(input())
lst = list(map(int, input().split()))
NGF = [-1] * N
counts = Counter(lst) # Counter 라이브러리 사용

idx = -1
indexStack = [] # 오등큰수를 찾지 못하고 있는 값들의 인덱스 리스트
for num in lst:

    idx += 1
    while indexStack and counts[lst[indexStack[-1]]] < counts[num]:
        NGF[indexStack.pop()] = num

    indexStack.append(idx)

print(*NGF)

 

[풀이]

여기에서는 '수열에서 등장한 횟수'가 기준이므로, 우선 Counter 라이브러리를 사용해서 입력받은 수열에서, 각 value가 등장하는 횟수를 체크한다. 그리고 stack의 top 원소의 숫자와, 현재 숫자의 등장 횟수를 비교하여 오등큰수 조건을 만족한 경우 조건이 더 이상 충족하지 않을 때까지 pop()을 진행하며 오등큰수를 찾아주도록 한다.

 

위에서 풀었던 오큰수 문제와 비슷해, 1~2주가 지나고 stack을 잘 이해했는지 복습 겸 풀기에 적당한 듯 하다.

 

 

---

 

2) BOJ - 3015. 오아시스 재결합 (플레 5)

 

[문제]

줄에 서 있는 사람의 키가 주어졌을 때, 서로 볼 수 있는 쌍의 수를 구한다.

이 때, 서로 볼 수 있는 조건은 A B 두 사람 사이에 A 또는 B보다 키가 큰 사람이 없는 경우이다.

 

[코드]

""" 동일한 키가 없을 경우의 알고리즘 """

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

stack = []
result = 0

# num을 기준으로 이전 원소 중 조건에 부합한 원소들만을 stack에서 관리하여, 만족하는 쌍을 찾도록 한다.
for num in lst:

    # stack에 num을 기준으로 높은 수만 있다. (내림차순)
    while stack and stack[-1] < num:
        stack.pop()
        result += 1

    # stack은 내림차순으로 정렬되어 있으므로, num에서 볼 수 있는 원소는 stack[-1] 밖에 없다.
    if stack and stack[-1] > num:
        result += 1

    stack.append(num)

print(result)

 

""" 동일한 키가 있을 경우의 알고리즘 """

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

stack = []
result = 0

# num을 기준으로 이전 원소 중 조건에 부합하 원소들만을 stack에서 관리하여, 만족하는 쌍을 찾도록 한다.
for num in lst:

    # stack에 num을 기준으로 높은 수만 있다. (내림차순)
    while stack and stack[-1][0] < num:
        result += stack.pop()[1]


    if stack:
        if stack[-1][0] == num:
            cnt = stack.pop()[1]
            result += cnt

            # 동일한 원소를 pop() 했음에도, 여전히 stack에 원소가 남아있으면 이는 num보다 큰 수로 바라볼 수 있는 수이다.
            if stack:
                result += 1

            stack.append((num, cnt + 1))


        # stack은 내림차순으로 정렬되어 있으므로, num에서 볼 수 있는 원소는 stack[-1] 밖에 없다.
        elif stack[-1][0] > num:
            result += 1
            stack.append((num, 1))


    # 앞에서 조건에 맞추어 stack에 값을 넣어주므로, stack이 비어 있는 경우에만 넣도록 한다.
    else:
        stack.append((num, 1))

print(result)

 

[풀이]

"A 또는 B보다 키가 큰 사람이 없어야 한다."가 문제에서 요구하는 조건이다.

여기에서 A 또는 B와 동일한 키가 없을 경우와 있을 경우를 고려했을 때 문제의 복잡도가 달라진다.

 

---

 

 

동일한 키가 없을 경우.

stack.top < num 일 때, num 이후의 원소들은 num에 가로막혀 stack의 원소를 더 이상 바라보지 못한다.

(num 이후의 원소가 num보다 크든 작든 num에 막힌다.)

따라서 반복적으로 stack.top을 pop() 하여, stack에 남아 있는 원소들은 num보다 큰 원소들만 남도록 한다.

(이후 남아있는 stack의 원소들은 num보다 큰 수들만 내림차순으로 남아, num 이후의 원소가 num보다 클 경우 바라볼 수 있음)

 

이로 인해 stack에는 내림차순 정렬을 보장받게 되고, num은 pop()을 하면서 '서로 볼 수 있는 쌍'을 카운트 한다.

 

stack.top > num 일 때 또한, 내림차순 정렬로 인해 num에서 볼 수 있는 원소는 stack.top 밖에 존재하지 않는다.

 

 

 

동일한 키가 있을 경우. 신경써야 할 부분이 늘어난다.

우선 stack에 num Value만 넣는 것이 아니라, 본인과 동일한 수가 몇 개 있는지 체크할 수 있는 Value도 넣어주어야 한다.

 

아래 코드를 보자.

stack 값과 num이 동일할 때, 현재 num에서는 이전까지의 모든 동일한 수를 바라볼 있기 때문에 cnt 횟수를 모두 카운트 해주어야 한다. 그리고 여기에서 주의할 점이 앞서 했던 "stack[-1] > num" 과 같은 처리를 한 번 더 해주어야 한다. 이는 동일한 원소를 pop()을 했음에도, 여전히 stack에 원소가 남아있다면 stack의 특성 상 num보다 큰 수이므로 바라볼 수 있게 된다.

if stack[-1][0] == num:
    cnt = stack.pop()[1]
    result += cnt

    if stack:
        result += 1

    stack.append((num, cnt + 1))

 

나머지는 '동일한 키가 없을 경우'의 알고리즘과 유사하기 때문에 자세한 설명은 생략하고, 주석만 남기도록 한다.

 

 

---

 

 

[스터디 코드]

def solution(heights):
    stack = []
    result = 0

    for h in heights:
        count = 1
        while stack and stack[-1][0] <= h:
            height, height_count = stack.pop()
            result += height_count
            if height == h:
                count += height_count

        if stack:
            result += 1

        stack.append((h, count))

    return result

 

위 문제를 갖고 스터디를 진행했는데, 스터디원 분의 깔끔한 코드를 보고 소개드립니다.

 

stack[-1][0] < h 과 stack[-1][0] == h 일 때,

모두 stack에서 pop()을 하는 로직이므로, 동일한 키에 대해서만 추가로 count를 더하면 코드가 더 깔끔해질 수 있다.

그리고 while pop() 했음에도, 여전히 stack에 원소가 남아있으면 이는 num보다 큰 수로 바라볼 수 있는 수이기 때문에, +1을 하면 된다.

 

stack.append() 로직에서도 그렇다.

이전 풀이에서는 조건에 따라 (현재 키, 1) / (현재 키, 동일한 키 + 1) 를 진행했었다.

본 풀이에서는 기본 count를 1로 두고, 동일한 키에 대해서는 깔끔히 count += height_count 를 하면 훨씬 깔끔하게 작성될 수 있다.