[1920] 수 찾기

 

백준 Silver IV | 1920 | Python | 문제 링크

문제 설명

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 안에 존재하는지 알아내면 된다. 모든 정수의 범위는 -2^31 보다 크거나 같고 2^31 - 1 보다 작거나 같다.

출력

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

입출력 예

입력	출력
5 4 1 5 2 3 5 1 3 7 9 5	1 1 0 0 1

나의 풀이

풀이 1 - set 활용과 버그 수정

수를 저장해두고 각 질의마다 존재 여부를 확인하면 된다. 리스트에 저장하고 in 연산으로 확인하는 것이 가장 먼저 떠올랐다.

그런데 잠깐 멈췄다. N이 최대 100,000이고 M도 최대 100,000이다. 리스트에서 in 연산은 처음부터 끝까지 순서대로 찾아가는 O(n)이다. M개의 질의마다 N개를 탐색하면 최악의 경우 O(NM), 즉 10^10 연산이다. 시간 초과가 날 것 같았다.

set을 쓰면 된다. 해시 기반이라 in 연산이 O(1)로 고정된다. N개를 set에 넣는 데 O(N), M개 질의마다 O(1)로 확인하면 전체 O(N+M)이다. 여기까지는 깔끔하게 정리됐다고 생각했다.

import sys
input = sys.stdin.readline
n = int(input())
num_set = set(map(int, input().split()))
m = int(input())
arr = map(int, input().split())
for i in arr:
    if i in set:
        print(1)
    else:
        print(0)

제출하고 틀렸다는 결과가 떴다.

코드를 천천히 다시 읽었다.

if i in set:

이 줄이 눈에 걸렸다. set. 제가 만든 변수 이름은 num_set이었다.

set은 Python 내장 타입 자체의 이름이다. 변수가 아니다. i in set이라고 쓰면 정수 i가 set 타입 안에 있는지를 묻는 셈인데, 그것이 말이 되지 않으니 항상 False를 반환하게 된다. 모든 질의에 0을 출력하고 있었던 것이다.

코드를 작성할 때 set을 직접 타이핑하면서 변수 이름 num_set을 가리키는 것처럼 읽었다. 머릿속으로 자동으로 채워 넣었던 것 같다.

Python에는 이런 함정이 꽤 많다. list, dict, int, str, type, id, input. 모두 내장 이름이다. 이것들을 변수명으로 쓰는 순간 원래 의미를 덮어씌우게 된다.

버그를 고치고 나서 다른 사람 풀이를 찾아봤다. 제 수정 코드와 다른 점이 하나 있었다. 질의마다 print()를 호출하는 대신 결과를 리스트에 모아두고 마지막에 한 번에 출력하는 방식이었다. M이 100,000이면 print()를 100,000번 호출하는 것인데, 그만큼 I/O가 발생한다.

풀이 2 - 최종 제출 코드

import sys
input = sys.stdin.readline

n = int(input())
num_set = set(map(int, input().split()))
m = int(input())
arr = map(int, input().split())

result = []
for i in arr:
    if i in num_set:
        result.append(1)
    else:
        result.append(0)

print("\n".join(map(str, result)))

풀이 3 - 제너레이터 표현식

import sys
input = sys.stdin.readline

n = int(input())
num_set = set(map(int, input().split()))
m = int(input())

print("\n".join("1" if i in num_set else "0" for i in map(int, input().split())))

루프 전체가 join의 인자로 들어간다. result 리스트도 없고, 변수도 적다. 흐름은 같다.

O(NM)이 느린 이유가 머리로는 알고 있었다. 직접 set으로 고치면서 O(N+M)으로 줄어드는 것이 실감이 됐다. 리스트를 쓰면 질의 하나마다 최대 N번 순회하고, set을 쓰면 질의 하나가 상수 시간에 끝난다. 같은 in 연산인데 뒤에 붙는 자료구조가 뭐냐에 따라 완전히 달라지는 것이다.