[백준][Python] 10816 숫자 카드 2

문제 링크

https://www.acmicpc.net/problem/10816

 

문제

숫자 카드는 정수 하나가 적혀져 있는 카드이다. 상근이는 숫자 카드 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           예제 출력1

10 6 3 2 10 10 10 -10 -10 7 3 8 10 9 -5 2 3 4 5 -10

3 0 0 1 2 0 0 2

 

---

풀이

import sys
input = sys.stdin.readline

def find_first(arr, target):
    left, right = 0, len(arr) - 1
    result = -1
    
    while left <= right:
        mid = (left + right) // 2
        
        if arr[mid] == target:
            result = mid
            right = mid - 1
        elif arr[mid] < target:
            left = mid + 1
        else:
            right = mid - 1

    if result != -1 :
        return result - 1
    else :
        return result

def find_last(arr, target):
    left, right = 0, len(arr) - 1
    result = -1
    
    while left <= right:
        mid = (left + right) // 2
        
        if arr[mid] == target:
            result = mid
            left = mid + 1
        elif arr[mid] < target:
            left = mid + 1
        else:
            right = mid - 1

    return result


n = int(input())
nums = list(map(int, input().split()))
nums.sort()

m = int(input())
targets = list(map(int, input().split()))
for target in targets :
    print(find_last(nums, target) - find_first(nums, target), end = ' ')

우선, 이분탐색을 사용하기 위해 nums에 입력한 숫자들을 정렬해 주었다.

그 뒤 target의 갯수을 찾기 위해 target의 가장 큰 인덱스를 구하는 함수와, target의 가장 작은 인덱스를 반환하는 함수를 호출해 주었다. 맨 처음에는 upper_bound와 lower_bound 함수를 생각하였으나, upper_bound는 초과하는 지점을 반환해 주어 만약 [1, 10] 중 타겟이 10이라면 오류가 발생할 수 있다. 따라서 약간의 변주를 주어 find_last(가장 큰 인덱스 반환) 함수와 find_first(가장 작은 인덱스 반환) 함수를 만들었다.

갯수를 구하기 위해선 (큰 인덱스) - (작은 인덱스) + 1 식이 필요하다. 하지만 무작정 메인에서 출력할 때 +1을 추가해주면 0번(target이 없음)이 1로 출력되는 일이 발생하기 때문에,  가장 작은 인덱스를 구할 때 target을 찾으면 -1 연산한 값을 반환하도록 하였다.