본문 바로가기

Algorithm Problems/이진 탐색

[백준/python] 7795번: 먹을 것인가 먹힐 것인가

문제

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

 

7795번: 먹을 것인가 먹힐 것인가

심해에는 두 종류의 생명체 A와 B가 존재한다. A는 B를 먹는다. A는 자기보다 크기가 작은 먹이만 먹을 수 있다. 예를 들어, A의 크기가 {8, 1, 7, 3, 1}이고, B의 크기가 {3, 6, 1}인 경우에 A가 B를 먹을

www.acmicpc.net


문제 요약

테스트 케이스개수 T를 입력 받고 A 집합, B 집합, A 집합의 크기, B 집합의 크기를 입력 받는다. 

 

A 집합과 B 집합을 비교하여 A가 더 큰 쌍이 몇 개가 있는지 출력한다.


코드

def binary_search(a, B): # a보다 작은 마지막 인덱스 요소 반환
    start = 0
    end = len(B) - 1
    res = -1

    while start <= end:
        mid = (start + end) // 2
        if B[mid] < a: 
        # B 리스트의 탐색 범위 중간 요소보다 a가 크면
            res = mid # res를 현 탐색 범위의 중간 인덱스로 저장
            start = mid + 1 # 탐색 범위를 mid + 1 ~ end 까지 절반으로 좁힘
        else:
            end = mid - 1 # 탐색 범위를 start ~ mid -1 까지 절반으로 좁힘
    return res

def main():
    T = int(input())

    for i in range(T):
        a_num, b_num = map(int, input().rstrip().split())
        A = list(map(int, input().rstrip().split()))
        B = list(map(int, input().rstrip().split()))

        B = sorted(B)
        A = sorted(A)

        count = 0

        for a in A:
            count += binary_search(a, B) + 1

        print(count)

if __name__ == "__main__":
    main()

코드 설명

단순히 이중 반복분을 사용하여 A와 B를 비교한다면 시간 초과 문제가 발생한다.

 

따라서 이진 탐색을 이용하여 문제를 풀 수 있다.

 

+ A 집합의 요소를 a, B 집합의 요소를 b라고 가정한다.

 

binary_search 함수는 A 집합 안에 각 요소 int형 a와 리스트 B를 입력받는다.

그리고 B에서 a보다 작은 인덱스 중 가장 큰 인덱스를 반환한다.


< 메인 함수 >

 

1. 먼저 입력 받은 A, B를 sorted함수를 이용하여 오름차순 정렬한다.

 

2. 반복문을 이용하여 A의 요소 a 하나씩 접근하고 binary_search 함수에 인자로 전달한다.

 

3. count 변수에 반환된 인덱스에 1을 더해 저장한다. 

 

+ B에서 a보다 작은 인덱스 중 가장 큰 인덱스 + 1 = B에서 a보다 작은 요소 개수


< biniary_search 함수>

 

1. start를 0, end를 B 길이 - 1로 설정한다.

 + res는 -1로 설정하여 탐색했을 때 a보다 작은 요소가 없다면 메인 함수에서 1을 더하여 a보다 작은 요소 개수가 0개라고 출력한다.

 

2. start와 end를 바꿔가며 B[mid]를 a와 비교한다.

+ mid는 (start + end) // 2 이므로 start와 end가 바뀔 때마다 (범위가 바뀔 때마다) mid도 바뀐다.

 

3. B[mid]가 a보다 작다면 res를 mid로 저장하고 start를 mid + 1로 변경하여 탐색 범위를 mid + 1 ~ end 까지좁힌다. 

+ 현재 res는 B에서 a보다 가장 작은 인덱스에 가까운 인덱스이다.

 

4. B[mid]가 a보다 같거나 크다면 end를 mid -1로 변경하여 탐색 범위를 start ~ mid - 1까지 절반으로 좁힌다. 

 

5. 반복문은 start가 end보다 커지면 종료되고 res를 반환한다.