백준 4949. 균형잡힌 세상(Python)

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

4949번: 균형잡힌 세상

 

이 문제의 조건은 다음과 같다.

1. 문자열의 괄호는 균형이 맞아야 한다.

2. 괄호는 일반적인 공식에 맞게 해야 한다.

3. 입력의 종료조건으로 맨 마지막에 온점 하나(".")가 들어온다.

 

그래서 이 3개의 조건에 초점을 두고 풀었다.

 

이번 문제의 경우에는 명확히 스택을 사용하는 문제이다. 문장에 있는 글자들을 하나씩 꺼내서 체크해주는 방법을 고민해보았는데, 2번 조건이 너무 거슬렸다.

2번 조건을 어떻게 할지가 너무 고민이었어서 우선은 스택의 특징을 살려서 다음과 같이 코드를 짰다.

while(True):
    word=input()

    stack=[]
    if word=='.':
        break

    for w in word:
        if w=='(' or w=='[':
            stack.append(w)
        elif w==')': # ) 인 경우
            if len(stack)!=0 and stack[-1]=='(':
                stack.pop()
            else:
                stack.append(w)
                break
        elif w==']':
            if len(stack)!=0 and stack[-1]=='[':
                stack.pop()
            else:
                stack.append(w)
                break

    if len(stack)==0:
        print('yes')
    else:
        print('no')

 

정답은 맞았는데 시간이 356ms가 나왔다. Python 1등 기록이 52ms인 걸 생각하면 참 터무니없이 과한 시간이다.

 

그래서 코드를 리팩토링했는데, 아래는 그 리팩토링한 코드이다.

import sys

def cal(target):
    stack = []
    for s in target:
        if not len(stack) and s in [")", "]"]:
            return "no"
        elif s in ["(", "["]:
            stack.append(s)
        elif s == ")":
            if stack[-1] != "(":
                return "no"
            stack.pop()
        elif s == "]":
            if stack[-1] != "[":
                return "no"
            stack.pop()
    if not len(stack):
        return "yes"
    return "no"

while True:
    target = sys.stdin.readline().rstrip()
    if target == ".":
        break
    print(cal(target))

 

84ms까지 줄일 수 있었다.

 

스택에서 괄호를 지울 수 없는 경우가 생길 때 바로 "no"를 반환하게끔 했기 때문에 불필요한 식을 줄일 수 있었다.