백준 1920: ArrayList 검색을 HashSet으로 바꾸기

백준 1920번 수 찾기는 주어진 수가 앞서 입력한 목록에 있는지 확인하는 문제입니다. 있으면 1, 없으면 0을 출력합니다.

목록을 하나씩 찾았더니 시간 초과가 났습니다

처음에는 입력을 ArrayList에 넣고 contains()로 존재 여부를 확인했는데 시간 초과가 났습니다. 검색할 때마다 목록을 앞에서부터 훑으므로 원소가 N개라면 한 번의 검색에 최대 O(N)이 필요합니다.

당시 실패한 코드는 다음과 같습니다.

Java

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

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));

        String N = br.readLine();
        String[] A = br.readLine().split(" ");
        String M = br.readLine();
        String[] strings = br.readLine().split(" ");

        List<String> arrayList = new ArrayList<>(Arrays.asList(A));

        for (String s : strings) {
            if (arrayList.contains(s)) {
                bw.write(String.valueOf(1));
            } else {
                bw.write(String.valueOf(0));
            }
            bw.newLine();
        }
        bw.flush();
    }
}

존재 여부만 필요해서 HashSet으로 바꿨습니다

이 문제에서는 저장 순서가 필요하지 않아서 목록을 HashSet으로 바꾸고 같은 방식으로 contains()를 호출했습니다. 입력값은 숫자로 변환하지 않고 문자열로 보관했습니다.

당시에는 HashSet으로 시간 복잡도를 O(log N)으로 줄였다고 적었지만 정확하지 않았습니다. Java 17 HashSet 문서는 해시가 원소를 적절히 분산한다는 조건에서 contains() 등 기본 연산이 상수 시간이라고 설명합니다. 여기서는 검색 한 번을 평균 O(1)로 보는 것이 맞습니다.

Java

import java.io.*;
import java.util.HashSet;
import java.util.List;

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));

        String N = br.readLine();
        String[] A = br.readLine().split(" ");
        String M = br.readLine();
        String[] strings = br.readLine().split(" ");

        HashSet hashSet = new HashSet<>(List.of(A));
        for (String s : strings) {
            if (hashSet.contains(s)) {
                bw.write(String.valueOf(1));
            } else {
                bw.write(String.valueOf(0));
            }
            bw.newLine();
        }
        bw.flush();
    }
}