2776번: 암기왕

문제

연종이는 엄청난 기억력을 가진 '암기왕'입니다. 동규는 연종이가 하루 동안 본 정수들을 모두 기억하는지 테스트하려고 합니다. 동규가 부르는 수들이 연종이가 수첩 1에 적어 놓은 하루 동안 본 수들 중에 있는지 없는지 판별하는 프로그램을 작성하는 문제입니다.

• 입력: 첫째 줄에 테스트 케이스의 개수 T가 주어집니다. 각 테스트 케이스의 첫째 줄에는 수첩 1에 적어 놓은 정수의 개수 N(1 ≤ N ≤ 1,000,000)이 주어집니다. 둘째 줄에는 수첩 1에 적혀 있는 정수들이 주어집니다. 셋째 줄에는 수첩 2에 적어 놓은 정수의 개수 M(1 ≤ M ≤ 1,000,000)이 주어집니다. 넷째 줄에는 수첩 2에 적혀 있는 정수들이 주어집니다.
• 출력: 수첩 2에 적혀있는 순서대로, 해당 정수가 수첩 1에 있으면 1을, 없으면 0을 출력합니다.

링크: https://www.acmicpc.net/problem/2776

풀이: HashSet을 활용한 빠른 탐색

수첩 1과 수첩 2에 적힌 정수의 개수가 최대 1,000,000개이므로, 이중 반복문으로 단순 탐색(O(N*M))을 수행하면 시간 초과가 발생합니다. 이를 해결하기 위해 O(1)의 시간 복잡도로 탐색이 가능한 HashSet을 사용합니다. 정렬과 이분 탐색(O(M log N))을 사용하는 방법도 있지만, 자바의 해시 자료구조를 활용하면 코드가 훨씬 간결해집니다.

• 데이터 저장: 수첩 1에 적혀 있는 N개의 정수를 입력받으면서 모두 HashSet(set)에 추가(add)합니다.
• 포함 여부 확인: 수첩 2에 적혀 있는 M개의 정수를 순서대로 입력받습니다. 각각의 숫자가 set에 존재하는지 contains() 메서드를 통해 확인합니다.
• 결과 출력: contains() 결과가 참(true)이면 수첩 1에 존재하는 것이므로 1을, 거짓(false)이면 존재하지 않는 것이므로 0을 BufferedWriter를 이용해 출력합니다.

코드

import java.io.*;
import java.util.*;

public class Main {
    public static void main(String[] args) throws IOException {
       BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
       BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

       int t = Integer.parseInt(br.readLine());

       while (t-- > 0) {
          int n = Integer.parseInt(br.readLine());
          StringTokenizer st = new StringTokenizer(br.readLine());
          Set<Integer> set = new HashSet<>();
          for (int i = 0; i < n; i++) {
             set.add(Integer.parseInt(st.nextToken()));
          }

          int m = Integer.parseInt(br.readLine());
          st = new StringTokenizer(br.readLine());
          for (int i = 0; i < m; i++) {
             int x = Integer.parseInt(st.nextToken());
             if (set.contains(x)) {
                bw.write("1\n");
             } else {
                bw.write("0\n");
             }
          }
       }

       bw.flush();
    }
}

알고리즘 분류

자료 구조 (Data Structures), 해시를 사용한 집합과 맵 (Hash Set/Map), 이분 탐색 (Binary Search)