정렬된 배열에서 특정 수의 개수 구하기

📌문제

N개의 원소를 포함하고 있는 수열이 오름차순으로 정렬되어 있다. 이때 이 수열에서 x가 등장하는 횟수를 계산하시오.

단 이 문제는 시간 복잡도 O(logN)으로 알고리즘을 설계하지 않으면 시간 초과 판정을 받습니다.

 

📌풀이

모든 원소가 정렬이 된 상태로 입력되므로 이진 탐색을 이용하여 값이 x인 원소의 개수를 O(logN)에 찾아낼 수 있음

수열 내 x가 존재한다면 연속적으로 나열되어 있음

따라서  x가 처음 등장한 인덱스와 마지막으로 등장한 인덱스의 차를 계산

 

이진 탐색 함수 2개를 사용하여 문제 해결

1. 가장 첫번째 위치를 찾는 이진 탐색 함수

2. 가장 마지막 위치를 찾는 이진 탐색 함수

 

📌코드

# 정렬된 수열에서 값이 x인 원소 개수를 세는 메서드
def count_by_value(array, x):
    # 데이터 개수
    n = len(array)

    # x가 처음 등장한 인덱스
    a = first(array, x, 0, n - 1)

    # 수열에 x가 존재하지 않는 경우
    if a == None:
        return 0

    b = last(array, x, 0, n - 1)

    # 개수 반환
    return b - a + 1


# 처음 위치를 찾는 이진 탐색
def first(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2
    # 해당 값을 가지는 원소 중에서 가장 왼쪽에 있는 경우만 인덱스 반환
    if (mid == 0 or target > array[mid - 1]) and array[mid] == target:
        return mid
    # 중간점의 값보다 찾고자 하는 값이 작거나 같은 경우 왼쪽 확인
    elif array[mid] >= target:
        return first(array, target, start, mid - 1)
    # 중간점의 값보다 찾고자 하는 값이 큰 경우 오른쪽 확인
    else:
        return first(array, target, mid + 1, end)


# 마지막 위치 찾는 이진 탐색
def last(array, target, start, end):
    if start > end:
        return None
    mid = (start + end) // 2
    # 해당 값을 가지는 원소 중에서 가장 오른쪽에 있는 경우만 인덱스 반환
    if (mid == n - 1 or target < array[mid + 1]) and array[mid] == target:
        return mid
    # 중간점의 값보다 찾고자 하는 값이 작은 경우 왼쪽 확인
    elif array[mid] > target:
        return last(array, target, start, mid - 1)
    # 중간점의 값보다 찾고자 하는 값이 큰 경우 오른쪽 확인
    else:
        return last(array, target, mid + 1, end)


n, x = map(int, input().split())
array = list(map(int, input().split()))

# 값이 x인 데이터 개수 계산
count = count_by_value(array, x)

if count == 0:
    print(-1)
else:
    print(count)

 

📌참고

https://g.co/kgs/eyd5SSd

이것이 취업을 위한 코딩 테스트다 with 파이썬