💻 백준 10799번: 쇠막대기

문제 소개 🧐

쇠막대기와 레이저의 배치가 괄호 문자열로 주어진다.

  • 쇠막대기는 (와 )로 표현된다.
  • 레이저는 인접한 () 쌍으로 표현된다.

레이저가 쇠막대기를 자를 때, 잘려진 쇠막대기 조각의 총 개수를 구하는 문제다.

입력 📝

  • 쇠막대기와 레이저의 배치를 나타내는 괄호 표현 (최대 100,000자)

출력 📤

  • 잘려진 조각의 총 개수

나의 풀이 ✨

괄호와 관련된 문제이므로 스택 자료구조를 활용하는 아이디어를 떠올렸다.

Idea 1: 레이저와 쇠막대기 분리 후 계산

첫 번째 아이디어는 먼저 전체 문자열을 순회하여 레이저의 위치와 쇠막대기의 위치(시작, 끝)를 각각 리스트에 저장하는 것이었다. 그 후, 각 쇠막대기마다 모든 레이저를 확인하며 내부에 포함되는 레이저의 개수를 세어 조각 수를 계산했다.

  • 쇠막대기 조각 수 = (내부 레이저 수) + 1
Code
import sys

S = list(sys.stdin.readline().rstrip())
stack = []
lasers = []
irons = []

for idx, elm in enumerate(S):
    stack.append((idx, elm))
    if elm == ')':
        cur_idx, cur_elm = stack.pop()
        prev_idx, prev_elm = stack.pop()
        if (cur_idx - prev_idx) == 1:
            lasers.append(prev_idx)
        else:
            irons.append((prev_idx, cur_idx))

answer = 0
for iron in irons:
    answer += sum(list(map(lambda x: 1 if iron[0] < x < iron[1] else 0, lasers)))+1

print(answer)
Result: 시간 초과

이 방법은 쇠막대기의 수와 레이저의 수를 각각 N, M이라고 할 때, 마지막 계산 과정에서 O(N*M)의 시간 복잡도를 가진다. 입력 크기가 최대 100,000이므로 당연히 시간 초과가 발생했다. 한 번의 순회로 모든 계산을 끝내는 방법이 필요했다.

Idea 2: 스택을 이용한 한 번의 순회

문자열을 한 번만 순회하면서 스택을 이용해 실시간으로 조각 수를 계산하는 방법이다.

  1. 문자열을 순회하며 (를 만나면 스택에 추가한다. (새로운 쇠막대기의 시작)
  2. )를 만나면 스택에서 하나를 꺼낸다.
    • 만약 바로 앞 문자가 (였다면, 이것은 레이저(())다. 이 레이저는 현재 스택에 쌓여있는 쇠막대기들을 모두 자르므로, answer에 현재 스택의 크기(len(stack))를 더해준다.
    • 만약 바로 앞 문자가 (가 아니었다면, 이것은 쇠막대기의 끝이다. 쇠막대기의 끝은 하나의 조각을 추가로 만들어내므로, answer에 1을 더해준다. (이 쇠막대기는 이미 이전의 레이저들에 의해 잘린 조각들이 계산된 상태이다.)
Code
import sys

S = list(sys.stdin.readline().rstrip())
stack = []

answer = 0
for idx, elm in enumerate(S):
    if elm == '(':
        stack.append((idx, elm))
    else:
        top_idx, top_val = stack.pop()
        if (idx - top_idx) == 1:  # 레이저
            answer += len(stack)
        else: # 쇠막대기의 끝
            answer += 1

print(answer)
Result: 성공! 🎉

이 방법은 문자열을 단 한 번만 순회하므로 O(N)의 시간 복잡도로 문제를 해결할 수 있다.

마무리 🤔

  • 스택의 활용: 괄호 쌍, 중첩 구조와 관련된 문제는 스택을 활용하면 효율적으로 해결할 수 있다는 것을 다시 한번 확인했다.
  • 효율적인 로직: 문제를 두 단계(탐색, 계산)로 나누는 대신, 한 번의 순회 과정에서 모든 계산을 처리하는 것이 얼마나 효율적인지 깨달았다.
  • 핵심 아이디어:
    • 레이저(())는 현재까지 열린 모든 쇠막대기를 자른다. → len(stack) 만큼 조각 추가
    • 쇠막대기의 끝())은 마지막 남은 한 조각을 의미한다. → 1 조각 추가

간단한 규칙을 통해 복잡한 문제를 해결하는 과정이 재미있었다.

⁂