[BOJ] 4949번: 균형잡힌 세상 / Python - dict 탐색으로 성능 올리기

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

입력으로 주어지는 각 문자열들에 대해, 소괄호 또는 대괄호의 열고 닫는 쌍이 갖추어져 있는지(="균형잡혀 있는지") 묻는 문제이다.

문제는 stack을 해결하면 풀린다.

문자열의 각 문자를 순회하면서,

• 여는 괄호일 때는 push하고
• 닫는 괄호일 때는 stack 내부에서 top의 위치에 매칭되는 여는 괄호가 이미 있는지 확인 -> 없으면 균형잡혀 있지 않은 것이다.

순회가 끝난 뒤, stack에 여는 괄호가 남아 있다면 이것도 균형잡혀 있지 않은 것이다.

관건은 구현이다.

 

먼저 아래 방법을 썼다.

p라는 문자열은 닫는 괄호에 대해 매칭되는 여는 괄호를 찾기 위한 저장소다. 이를 위해 여는 괄호 2개와 닫는 괄호 2개를 이어붙였다. 그래서 p 문자열 상에서 index() 계산을 사용해서 매칭되는 괄호를 찾는다.

정답 코드 1

import sys


def is_balanced(s):
    p = "([)]"
    stack = []
    for c in s:
        if c in p[0:2]:
            stack.append(c)
        elif c in p[2:]:
            if not stack or (stack[-1] != p[p.index(c) - 2]):
                return False
            stack.pop()
    return not stack


sys.stdout.write(
    "\n".join(
        ("yes" if is_balanced(s) else "no") for s in sys.stdin if s.rstrip() != "."
    )
)

문자열 p의 길이가 4밖에 안되기 때문에, p.index()가 linear search긴 해도 사실상 O(1)으로 동작할 것으로 생각했다. 

그렇지만, 실제로는 O(4)라고 볼 수 있다.

만약 입력 문장의 갯수가 매우 많아지면, 이 미세한 차이가 누적되서 성능에 영향을 끼칠 수 있다.

 

매칭되는 여는 괄호를 한번에 찾을 수 있는 방법을 사용하면 성능을 좀 더 최적화할 수 있다.

이를 위해, 해싱 기반으로 동작하는 dict를 사용했다.

정답 코드 2 - dict를 사용한 해싱 탐색

import sys


def is_balanced(s):
    stack = []
    pairs = {")": "(", "]": "["}
    for c in s:
        if c in "([":
            stack.append(c)
        elif c in ")]":
            if not stack or stack[-1] != pairs[c]:
                return False
            stack.pop()
    return not stack


sys.stdout.write(
    "\n".join(
        ("yes" if is_balanced(s) else "no") for s in sys.stdin if s.rstrip() != "."
    )
)

성능이 좀 더 개선된 것을 확인할 수 있다.