[Python/Silver4/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, bisect
input = sys.stdin.readline

class Solution:
    def find_number(self):
        N = int(input())
        A = list(map(int, input().split()))
        M = int(input())
        arr = list(map(int, input().split()))
        A.sort()
        
        for k in arr:
            lo, hi = 0, N-1
            find = False
            while lo <= hi:
                mid = lo + (hi-lo) // 2
                if A[mid] < k:
                    lo = mid + 1
                elif A[mid] > k:
                    hi = mid - 1
                else:
                    find = True
                    break
            if find:
                print(1)
            else:
                print(0)
    
    def find_num(self):
        N = int(input())
        A = list(map(int, input().split()))
        M = int(input())
        arr = list(map(int, input().split()))
        A.sort()
        for k in arr:
            idx = bisect.bisect_left(A,k)
            if idx >= N:
                print(0)
                continue
            if A[idx] == k:
                print(1)
            else:
                print(0)

if __name__ == "__main__":
    s = Solution()
    s.find_num()

 

 

|    비고

처음 작성했던 코드는 큐를 활용했다. 이 경우 시간 복잡도는 O(N*M)이 되는데, if a in A는 리스트에서 선형 탐색을 하기 때문에 O(N)이 걸리고, 이를 M번 반복하기 때문이다. 이 문제에서 N과 M은 둘 다 100,000까지 이므로 10 ^ 10 번 즉 연산 1억번을 넘어가기 때문에 시간 초과가 날 수밖에 없다. 따라서, 이 문제는 이분탐색(bisect) 혹은 set() 함수를 활용해서 풀면 된다.

import sys
from collections import deque
input = sys.stdin.readline

class Solution:
    def find_number(self):
        N = int(input())
        A = list(map(int, input().split()))
        M = int(input())
        arr = deque(map(int, input().split()))
        for i in range(M):
            a = arr.popleft()
            if a in A:
                print(1)
            else:
                print(0)

if __name__ == "__main__":
    s = Solution()
    s.find_number()