[Zoho 인터뷰] 특정 원소 개수 찾기

출처 : 나동빈 유튜브

 

풀이

---

python의 bisect를 이용하면 1분만에 풀 수 있지만, 이진탐색의 이해를 위해 이진탐색 메소드를 직접 구현해서 풀었다.

 

 

bisect 이용 풀이

from bisect import bisect_left, bisect_right


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

left = bisect_left(arr, x)
right = bisect_right(arr, x)

result = right - left
if result == 0:
    print("-1")
else:
    print(result)

 

이진탐색 메소드 직접 구현 풀이

def first(data, target, start, end):
    if start > end:
        return None

    mid = (start + end) // 2

    if data[mid] == target and (mid == 0 or data[mid-1] < target):
        return mid
    elif target <= data[mid]:
        return first(data, target, start, mid-1)
    else:
        return first(data, target, mid+1, end)


def last(data, target, start, end):
    if start > end:
        return None

    mid = (start + end) // 2
    if data[mid] == target and (mid == n-1 or data[mid+1] > target):
        return mid
    elif target < data[mid]:
        return last(data, target, start, mid-1)
    else:
        return last(data, target, mid+1, end)


def count_by_value(data, target):
    left = first(data, target, 0, n-1)
    if left == None:
        return 0

    right = last(data, target, 0, n-1)
    return right-left+1


n, x = map(int, input().split())
arr = list(map(int, input().split()))
result = count_by_value(arr, x)

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

 

설명

기존의 binary_search를 두개로 나누었다.

1. 하나는 특정원소중에 가장 낮은 인덱스를 구하는 메소드
2. 또 다른 하나는 특정원소중에 가장 높은 인덱스를 구하는 메소드

연속으로 나열되어있는 특정원소중에서 가장 낮은 인덱스를 구하기 위해서 기존 이진탐색 알고리즘과 달라지는 점이 있다.

원래의 이진탐색 알고리즘은 data[mid] == target 즉 해당 원소를 찾으면 바로 반환했었는데,

이 문제에서는 가장 낮은(높은) 인덱스를 찾아야한다.

 

• 기존 : if data[mid] == target
• 바뀜 : if data[mid] == target and (mid == 0 or data[mid-1] < target)

 

 

 (mid 가 0일때) 혹은 (mid 바로 앞의 원소보다 클 때) 특정 원소의 가장 낮은 인덱스를 찾을 수 있다.

그리고 target과 같지만 새로 추가된 조건을 통과하지 못했을 때 end 를 mid-1로 주어 왼쪽부분에서 다시 탐색을 시작하도록 만들면 된다.