[백준 7795] 먹을 것인가 먹힐 것인가 (Python)

www.acmicpc.net/problem/7795

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

 

문제풀이

N, M : 배열 A, B의 크기

A, B : 입력 배열

A_idx, B_idx : 투 포인터

count : A의 원소가 B의 원소보다 큰 경우의 갯수를 메모라이제이션

ans : 결과값

tmp : A_idx 번째의 A 원소에서, 현재 B_idx 부터 A[A_idx] > B[B_idx] 를 만족하는 B_idx 의 갯수

 

N, M 조건이 20000 밖에 되지 않아서 2중 for문으로 O(NM) 복잡도로 해결할 수도 있다. 하지만 이 방식은 상대적으로(?) 느리기 때문에 투 포인터와 메모라이제이션을 활용한다면 시간을 대폭 줄일 수 있다.

 

1. 배열 A와 B를 정렬해준다.

- 정렬을 하지 않으면 투 포인터를 정상적으로 사용할 수 없다.

 

2. A_idx와 B_idx를 이용하여 각각 값을 증가시켜 보면서 A의 원소가 B의 원소보다 큰지 확인한다.

- A의 원소가 B의 원소보다 큰 경우는 tmp, B_idx 증가

- A의 원소가 B의 원소보다 작거나 같은 경우는 A_idx - 1 번째 까지의 갯수 + tmp, A_idx 증가

- B의 마지막 원소까지 도달한 경우 남아있는 A 원소의 갯수는 A_idx - 1 번째의 갯수와 동일

 

3. 결과(sum(count))를 출력한다.

 

예제 테스트케이스에서 A, B 배열 정렬 후 메모라이제이션 된 결과를 간단하게 그림으로 그리면 다음과 같다.

 

자세한 설명은 코드 주석으로 설명되어 있습니다.

 

코드

import sys

# 조건 : A의 원소 중에서 B의 원소보다 큰 경우의 수 구하기

if __name__ == '__main__':
    for _ in range(int(input())):
        N, M = map(int, sys.stdin.readline().split())
        A = list(map(int, sys.stdin.readline().split()))
        B = list(map(int, sys.stdin.readline().split()))
        A.sort(); B.sort() # 정렬이 반드시 필요
        A_idx, B_idx = 0, 0 # 투 포인터
        count = [0] * N # A의 원소가 B의 원소보다 큰 경우의 갯수를 메모라이제이션
        ans = 0
        tmp = 0 # A_idx 번째의 A 원소에서, 현재 B_idx 부터 A[A_idx] > B[B_idx] 를 만족하는 B_idx 의 갯수
        while A_idx < N:
            if B_idx == M: # B 배열의 끝까지 도달한 경우
                # 3번 줄의 조건을 만족하는 A_idx 번째의 B_idx 갯수는 A_idx - 1 번째의 갯수 + tmp
                count[A_idx] = count[A_idx - 1] + tmp
                A_idx += 1
                tmp = 0 # tmp 초기화 => B 배열의 끝에 도달했기에 더 이상 tmp는 필요 없으므로 0 으로 초기화
                continue
            if A[A_idx] <= B[B_idx]: # B의 원소가 A의 원소보다 큰 경우
                if A_idx > 0: count[A_idx] = count[A_idx - 1] + tmp # A_idx - 1 번째 까지의 갯수 + tmp
                else: count[A_idx] += tmp # A_idx 가 0일 경우 예외 처리
                A_idx += 1 # A_idx 증가
                tmp = 0 # A_idx 증가 했으므로 tmp 초기화
                continue
            # A의 원소가 B의 원소보다 큰 경우
            tmp += 1
            if B_idx < M:
                B_idx += 1
        print(sum(count)) # 갯수 합 출력

 

결과

 

문제풀이나 코드에서 이상한 부분 지적은 언제나 환영합니다!