BOJ: 10815

Posted on In

문제

숫자 카드는 정수 하나가 적혀져 있는 카드이다. 상근이는 숫자 카드 N개를 가지고 있다. 정수 M개가 주어졌을 때, 이 수가 적혀있는 숫자 카드를 상근이가 가지고 있는지 아닌지를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 상근이가 가지고 있는 숫자 카드의 개수 N(1 ≤ N ≤ 500,000)이 주어진다. 둘째 줄에는 숫자 카드에 적혀있는 정수가 주어진다. 숫자 카드에 적혀있는 수는 -10,000,000보다 크거나 같고, 10,000,000보다 작거나 같다. 두 숫자 카드에 같은 수가 적혀있는 경우는 없다.
셋째 줄에는 M(1 ≤ M ≤ 500,000)이 주어진다. 넷째 줄에는 상근이가 가지고 있는 숫자 카드인지 아닌지를 구해야 할 M개의 정수가 주어지며, 이 수는 공백으로 구분되어져 있다. 이 수도 -10,000,000보다 크거나 같고, 10,000,000보다 작거나 같다.

출력

첫째 줄에 입력으로 주어진 M개의 수에 대해서, 각 수가 적힌 숫자 카드를 상근이가 가지고 있으면 1을, 아니면 0을 공백으로 구분해 출력한다.

Solution

Brute-force?

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
package main

import "fmt"

func main() {
	var N, M int
	fmt.Scan(&N)
	sg := make([]int, N)
	for i := 0; i < N; i++ {
		fmt.Scan(&sg[i])
	}
	fmt.Scan(&M)
	card := make([]int, M)
	res := make([]int, M)
	for i := 0; i < M; i++ {
		fmt.Scan(&card[i])
	}
	for i := 0; i < N; i++ {
		for j := 0; j < M; j++ {
			if sg[i] == card[j] {
				res[j] = 1
			}
		}
	}
	for i := 0; i < M; i++ {
		fmt.Print(res[i], " ")
	}
}

슬픈 예감은 역시 틀리지 않는,,,
너무 쉽게 풀려서 이상했지만 해당 문제의 N은 최대 500,000이고 M은 최대 10,000,000이다.
즉, 해당 알고리즘으로 문제를 풀 경우 시간 복잡도는 $O(N\times M)$으로 매우 커져 당연히 시간 초과가 발생한다.
그렇다면 Go에서 제공하는 자료형 중 하나인 map을 사용한다면?

Map?

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
package main

import "fmt"

func main() {
	var N, M int
	fmt.Scan(&N)
	var sg map[int]bool
	sg = make(map[int]bool)
	for i := 0; i < N; i++ {
		fmt.Scan(&M)
		sg[M] = true
	}
	fmt.Scan(&M)
	for i := 0; i < M; i++ {
		fmt.Scan(&N)
		_, flag := sg[N]
		if flag {
			fmt.Print(1, " ")
		} else {
			fmt.Print(0, " ")
		}
	}
}

map을 사용하여 상근이가 가지고 있는 숫자 카드를 sg 변수에 저장하고 주어진 숫자 카드와 비교할 때 sg 내에 키 값의 존재 여부를 flag 변수에 저장하고 문제 조건에 맞게 출력했다.
하지만 여전히 시간 초과가 발생하는데, 코드를 자세히 보면 N+M번의 입력이 존재하는데 이 또한 매우 크기 때문에 Go에서 빠른 입출력을 지원하는 bufio 패키지를 사용해야한다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
package main

import (
	"bufio"
	"fmt"
	"os"
)

func main() {
	var reader *bufio.Reader = bufio.NewReader(os.Stdin)
	var writer *bufio.Writer = bufio.NewWriter(os.Stdout)
	defer writer.Flush()

	var N, M int
	fmt.Fscan(reader, &N)
	var sg map[int]bool
	sg = make(map[int]bool)
	for i := 0; i < N; i++ {
		fmt.Fscan(reader, &M)
		sg[M] = true
	}
	fmt.Fscan(reader, &M)
	for i := 0; i < M; i++ {
		fmt.Fscan(reader, &N)
		_, flag := sg[N]
		if flag {
			fmt.Fprint(writer, "1 ")
		} else {
			fmt.Fprint(writer, "0 ")
		}
	}
}