[Algorithm] 백준 10799 쇠막대기 | 스택 (파이썬)

첫 번째 시도, 시간초과. 두 번째 약간의 힌트를 얻고 통과! 처음 생각하던 아이디어 그대로 풀었으면 시간 초과 안 나고 잘 풀었을 텐데 반복문을 두 번이나 쓰면서 시간초과로 실패했다. 그래도 문제가 아주 어렵지 않고, 마지막에 넣은 값을 빼는 스택을 활용한다는 아이디어만 바로 얻으면 금방 풀어낼 수 있다! 자 이제 풀어보자.

 

📝 백준 10799 쇠막대기

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

서로 마주 보는 ()는 레이저, 나머지 괄호 (--)는 막대기를 의미한다. 예시 이미지에서 첫 번째 괄호는 레이저지만, 앞 뒤로 괄호가 없어서 막대기가 없다는 의미이므로 자를 막대기가 없는 것! 이렇게 잘린 막대기의 수를 찾으면 된다.

 

💡 문제 풀이 아이디어

"("를 만나면 막대기의 시작을 의미한다. 그리고 ")" 만나면 가장 마지막에 시작한 막대기의 끝을 의미한다. '가장 마지막에 시작한' 에서 막대기를 저장하는 자료 구조로 스택을 활용할 수 있다는 아이디어를 얻을 수 있다. ( 를 만나서 막대기가 여전히 있음을 저장하고, 레이저를 만나면 살아있는 막대기의 개수를 세서 잘린 막대기 수를 셀 수 있다.

 

그림을 보면서 생각해봤을 땐, 레이저가 동작했을 때 남아있는 막대기의 수를 더하고, 막대기가 끝났을 때 남은 막대기 1을 더해주면 된다. (아이패드 애플 펜슬 인식이 이상해서 +가 이상하게 그려졌지만, 어쨌든 O를 그린 개수만큼 더해지는 것!)

 

📝 문제 풀이 (파이썬)

# 시간 복잡도 O(N)

arr = input()
stack = []
result = 0

for i, x in enumerate(arr):
	# ( 를 만나면 막대기를 시작
    if x == "(":
        stack.append("(")
    # ) 를 만났는데, 이전 막대기가 ( (레이저 O)
    elif arr[i - 1] == "(" and x == ")":
        stack.pop()  # stack에서 막대기를 빼고,
        result += len(stack)  # 지금까지 살아있는 막대기의 수를 더한다
    # ) 를 만나면 막대기 끝 (레이저 X)
    else:
        stack.pop()  # stack에서 막대기를 빼고,
        result += 1  # 남은 막대기 수 하나를 더한다

print(result)

 

여기서 몇 가지 보완, 수정을 한다면, enumerate로 사용하지 않고 이전값을 저장하는 변수를 하나 만들어도 될 것 같다. index가 필요한 이유는 이전 막대기가 무엇인지 알아야 하기 때문이니까! 그리고 elif로 () 레이저가 동작하는지 조건을 확인하는데, 그냥 "(" 조건 확인 후 else 구문에서 이전 값을 확인해도 될 수도?

 

🚨 문제 풀이 (시간 초과 났던 처음 문제 풀이)

# 시간 복잡도 O(N²)

arr = input()
stack = []
result = 0

for i, x in enumerate(arr):
    if x == "(":
        stack.append(1)
    elif x == ")" and arr[i - 1] == "(":
        stack = [s + 1 for s in stack]  # 여기서 반복문이 한 번 더 동작
        stack.pop()
    else:
        result += stack.pop()

print(result)

 

이 풀이는 각 막대기가 잘리는 수를 모두 stack에 담아두고 활용하는 방법이다. 레이저를 만나면 stack에 담긴 막대기 값을 1씩 더하는 반복문이 중간에 동작한다. 여기서 시간초과가 발생한 것 같다..