백준) 1920 - 파이썬

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

(실버 4 문제)

1920번: 수 찾기

입력

5
4 1 5 2 3
5
1 3 7 9 5

출력

1
1
0
0
1

내 코드

import sys
N=int(sys.stdin.readline())
data1=list(map(int,sys.stdin.readline().split()))
M=int(sys.stdin.readline())
data2=list(map(int,sys.stdin.readline().split()))

for i in range(M):
  result=data1.count(data2[i])
  print(result)

시간 초과가 난 코드이다.

처음에는 input()을 사용해서 시간 초과가 났다고 생각해서 import sys를 통해서 더 빠르게 입력을 받는 방식으로 고쳤으나 여전히 시간 초과였다.

for문에서 count를 이용해서 숫자가 몇 개 있는지 알아내는 방식으로 사용했다. 1 또는 0만 출력하면 되는 구조를 굳이 count를 사용할 필요는 없었다. 코드를 짤 당시에 count밖에 생각이 나지 않아서 count로 무작정 작성하였는데 이유없이 직관적으로 작성하는 방식을 고칠 필요가 있다.

풀이

import sys
N=int(sys.stdin.readline())
data1=set(map(int,sys.stdin.readline().split()))
M=int(sys.stdin.readline())
data2=list(map(int,sys.stdin.readline().split()))

for i in data2:
    if i in data1:
        print(1)
    else:
        print(0)

data1을 리스트에서 set으로 변경하니 바로 정답 처리가 되었다.

자료구조의 차이 때문인데 리스트 자료구조의 경우 search 시간이 O(n) 시간 걸린다.

반면 set 자료 구조의 경우 search 시간이 O(1) 시간이 걸린다.

즉 탐색 시간에 있어서 내 코드는 O(n^2) 만큼 걸리게 되는 것이고, 풀이는 O(n)만큼 걸리게 된다.

data1의 경우 data2의 탐색의 기준이 되기 때문에 중복되는 값이 없고, 그렇기에 set을 통해서 구현해도 상관이 없다. 굳이 리스트로 작성할 필요가 없었던 것이다. 

느낀 점

단순히 문제를 보고 기계적으로 작성하는 것이 아닌, 시간 초과라는 요소까지 생각해서 짜야한다는 것을 알게 되었다.

또한 자료 구조에 따라 탐색 시간이 바뀌어서 시간이 걸릴만한 요소를 줄일 수 있다는 사실을 알게 되었다.