백준) 11723 - 파이썬

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

(실버 5)

11723번: 집합

입력

26
add 1
add 2
check 1
check 2
check 3
remove 2
check 1
check 2
toggle 3
check 1
check 2
check 3
check 4
all
check 10
check 20
toggle 10
remove 20
check 10
check 20
empty
check 1
toggle 1
check 1
toggle 1
check 1

출력

1
1
0
1
0
1
0
1
0
1
1
0
0
0
1
0

내 풀이

import sys
input=sys.stdin.readline
n=int(input())
a=set()
for i in range(n):
    s=list(input().split())
    if s[0]=='add':
        if int(s[1]) in a:
            continue
        else:
            a.add(int(s[1]))
    elif s[0]=='remove':
        if int(s[1]) not in a:
            continue
        else:
            a.remove(int(s[1]))
    elif s[0]=='check':
        if int(s[1]) in a:
            print(1)
        else:
            print(0)
    elif s[0]=='toggle':
        if int(s[1]) in a:
            a.remove(int(s[1]))
        else:
            a.add(int(s[1]))
    elif s[0]=='all':
        a.clear()
        for j in range(1,21):
          a.add(j)
    elif s[0]=='empty':
        a.clear()

시간 초과가 난 코드이다

아무래도 all을 구현하는 과정에 있어서 for문을 한 번 더 사용해서 시간 복잡도가 O(n^2)가 되어서 시간 초과가 된거 같다.

또한 remove를 이용하지 않고 discard를 이용해서 보다 안전하게 제거가 가능하다.

개선 코드

import sys
input=sys.stdin.readline
n=int(input())
a=set()
for i in range(n):
    s=list(input().split())
    if s[0]=='add':
        if int(s[1]) in a:
            continue
        else:
            a.add(int(s[1]))
    elif s[0]=='remove':
        a.discard(int(s[1]))
    elif s[0]=='check':
        if int(s[1]) in a:
            print(1)
        else:
            print(0)
    elif s[0]=='toggle':
        if int(s[1]) in a:
            a.discard(int(s[1]))
        else:
            a.add(int(s[1]))
    elif s[0]=='all':
        a=set(range(1,21))
    elif s[0]=='empty':
        a.clear()

discard와 all 부분을 고쳐서 시간 초과를 해결한 코드이다. set의 경우 범위값을 지정시 range와 같이 결합해서 저장이 가능하다.

느낀 점

set 함수에 대해서 잘 모르는 부분이 많았는데 이번 문제를 통해 이해도가 높아졌다.