[백준] 1920번: 수 찾기 (Python)

[1920번] 수 찾기

문제

N개의 정수 A[1], A[2], …, A[N]이 주어져 있을 때, 이 안에 X라는 정수가 존재하는지 알아내는 프로그램을 작성하시오.

입력

첫째 줄에 자연수 N(1 ≤ N ≤ 100,000)이 주어진다. 다음 줄에는 N개의 정수 A[1], A[2], …, A[N]이 주어진다. 다음 줄에는 M(1 ≤ M ≤ 100,000)이 주어진다. 다음 줄에는 M개의 수들이 주어지는데, 이 수들이 A안에 존재하는지 알아내면 된다. 모든 정수의 범위는 -231 보다 크거나 같고 231보다 작다.

출력

M개의 줄에 답을 출력한다. 존재하면 1을, 존재하지 않으면 0을 출력한다.

---

실패한 풀이)

import sys
input = sys.stdin.readline

N = int(input())
a_list = list(map(int, input().split()))
M = int(input())
b_list = list(map(int, input().split()))
cnt_list = [0 for _ in range(M)]

for i in range(M):
    for a in a_list:
        if a == b_list[i]:
            cnt_list[i] += 1
            break
        
for i in cnt_list:
    print(i) 

문제에 시간 제한이 1초인데, for 문을 이중으로 사용하다 보니 시간복잡도가 O(MN)이 되어 시간초과가 발생하였다. 이진 탐색(binary search) 알고리즘을 사용하게 되면 O(MlogN)으로 시간복잡도가 훨씬 작기 때문에 이 방식으로 해결해야 한다.

풀이 1) 이진 탐색

import sys
input = sys.stdin.readline

N = int(input())
a_list = list(map(int, input().split()))
a_list.sort()

M = int(input())
b_list = list(map(int, input().split()))

for b in b_list:
    start = 0
    end = N - 1
    isIn = False

    while start <= end:
        mid = (start + end) // 2

        if b == a_list[mid]:
            isIn = True
            print(1)
            break
        elif b > a_list[mid]:
            start = mid + 1
        else:
            end = mid - 1
    
    if not isIn:
        print(0)

이진 탐색 알고리즘에 관련 내용은 여기 이진탐색 포스팅을 참고하면 된다.

풀이 2) 집합 set 자료형

import sys
input = sys.stdin.readline

N = int(input())
a_set = set(map(int, input().split()))

M = int(input())
b_list = list(map(int, input().split()))

for b in b_list:
    if b in a_list:
        print(1)
    else:
        print(0)

자료형 set을 사용하게 되면 사용하는 for 문이 하나로, 시간 복잡도가 줄고, 이진 탐색 방식보다도 더 간단하게 해결이 가능하다.