[Silver IV] 숫자 카드 2 - 10816

문제 링크

성능 요약

메모리: 224840 KB, 시간: 1088 ms

분류

이분 탐색(binary_search), 자료 구조(data_structures), 해시를 사용한 집합과 맵(hash_set), 정렬(sorting)

문제 설명

숫자 카드는 정수 하나가 적혀져 있는 카드이다. 상근이는 숫자 카드 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개의 수에 대해서, 각 수가 적힌 숫자 카드를 상근이가 몇 개 가지고 있는지를 공백으로 구분해 출력한다.

풀이 과정

문제는 단순하다.
카드들을 모두 입력받은 뒤, 입력한 숫자 값의 카드가 몇개 있는지 출력해주는 문제다.

숫자로 되어있는 배열 내에서 입력받은 값을 찾아야 하므로, 이 문제는 이분탐색 문제이다.

먼저 입력받은 카드들을 오름차순으로 정렬한다.

그리고 이분 탐색을 하면서 입력받은 값을 O(logN)의 속도로 찾아간다.
찾았다면 우선 1개가 있으므로, cnt변수에 1을 더 해준다.

찾았다면 while문을 통해 찾은 값의 왼쪽, 오른쪽에 같은 값이 있는지 찾는다.

찾았다면 cnt에 1씩 더해준다.
위 과정이 모두 끝났다면 cnt를 출력해준다.
여기서 하나도 찾지 못 했다면, 0이 출력될 것이다.

하지만 이렇게 이분 탐색'만' 해서 제출한다면 시간 초과가 뜬다.
왜냐하면 이 전에 이미 찾은 값을 또 찾아야 할 때, 같은 과정을 반복해야 하기 때문이다.

그래서 dict라는 dictionary 자료형을 만들고, 찾은 값을 key, cnt를 value로 한다.
그리고 새로운 값을 찾을 때, 이미 dict 안에 있으면 굳이 똑같은 값을 찾아주지 않아도 되므로 불필요한 과정이 생략된다!

코드

import sys

input = sys.stdin.readline

N = int(input())
A = sorted(list(map(int, input().split())))

M = int(input())
B = list(map(int, input().split()))

dict = {}

for b in B :
    if b not in dict :
        cnt = 0
        start = 0
        last = N - 1
        while start <= last :
            mid = (start + last) // 2
            if b > A[mid] :
                start = mid + 1
            elif b < A[mid] :
                last = mid - 1
            else :
                cnt += 1
                # 왼쪽에 있는지 찾음
                prev = -1
                while mid + prev >= 0 and b == A[mid + prev] :
                    cnt += 1
                    prev -= 1
                    
                # 오른쪽에 있는지 찾음
                next = 1
                while mid + next < N and b == A[mid + next] :
                    cnt += 1
                    next += 1
                break
        dict[b] = cnt
    print(dict[b], end=" ")