[Code Tree][Python] 앞에서부터 삭제하기 2

Overview

• 체감 난이도: ★★★★☆ (효율성 측면에서 시간 복잡도 줄이는 부분이 어려웠다)
• 문제 유형: 우선순위큐/힙
• 풀이 상태: 답안참고 / 스스로 해결
• 추후: 다시 풀어보기 / 간단 복습 / 완벽 이해 

---

[문제]

https://www.codetree.ai/missions/8/problems/delete-it-from-the-beginning-2/description

[코드]

(시간초과 코드)

import heapq, sys
input = sys.stdin.readline

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

tmp, mx = [], 0
for k in range(1, n-1):  # k = 1,2,3
    lst = arr[k:]

    if len(lst) > 1:
        heapq.heapify(lst)
        x = heapq.heappop(lst)
        if mx < sum(lst)/len(lst):
            mx = sum(lst)/len(lst)

print("{:.2f}".format(mx))

 

• for k in range(1, n-1): O(n)
• lst = arr[k:]: O(n-k) → O(n-1) * O(n-2) ... → O(n^2) 가까움 
• heapq.heapify(lst) : O(n log n)
• heapq.heappop(lst) : O(log n)
• sum(lst)/len(lst) : O(n)
• 종합 시간 복잡도: O(n^2 log n)

 

(정답 코드)

1. 우선순위큐(힙)으로 풀이 

Idea ) 주어진 값들의 뒤에서부터 계산하며 답을 구한다. (k = N-2, N-3, ... , 1)

 

→ K = 1,2,3,4, N-2까지 순서대로 K개의 수를 지우고 난 뒤 남아있는 수들 중 최솟값을 찾는 시간 ▶︎ O(N^2)

→ 만약 주어진 수들을 뒤에서부터 계산하여 앞으로 이동한다면 지금까지 거쳐온 수들 중 최솟값을 찾는 것은 Priority Queue에 원소를 넣어 그 중 최솟값을 찾는 것과 같으므로 계속 이전 Priority Queue에 들어 있던 값을 재활용할 수 있다. 즉, K = N - 2 일때부터 답을 내고, 다시 원소를 하나 추가하여 최솟값을 구하여 K = N - 3 일때의 답을 내는 식으로 K = 1 일때까지 진행하면 된다. 

 

import heapq

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

sum_val = 0  # 현재 남은 원소들의 합을 저장하는 변수
max_avg = 0  # 계산된 최대 평균을 저장하는 변수 
pq = []  # 남아있는 원소들을 삽입하여 최솟값을 찾는 리스트 (우선순위 큐 리스트)

heapq.heappush(pq, arr[n - 1])  # 초기: 가장 마지막 원소 추가 
sum_val += arr[n - 1]

# k가 N - 2일 때부터 1일 때까지 거꾸로 탐색
# priority queue를 이용하여 진행
for i in range(n - 2, 0, -1):
    # 앞에서부터 K개를 삭제하고 나면 뒤에 i ~ n - 1 까지의 숫자만이 남는다.
    heapq.heappush(pq, arr[i])
    sum_val += arr[i]

    # 남아있는 정수 중 가장 작은 숫자를 찾아
    # 그 숫자를 제외한 평균을 구한다
    min_num = pq[0]
    avg = (sum_val - min_num) / (n - i - 1)

    # 평균이 최대가 된다면 정답을 현재 평균으로 갱신
    if max_avg < avg:
        max_avg = avg

print(f"{max_avg:.2f}")

 

2. 누적합으로 풀이 

n = int(input())
arr = [0] + list(map(int, input().split()))

# 왼쪽부터 누적합(arr의 첫 번째 요소부터 i번째 요소까지)을 계산하여 저장하는 배열 => O(N)
prefix_sum = [0] * (n + 1)
for i in range(1, n + 1):
    prefix_sum[i] = prefix_sum[i - 1] + arr[i]

# 오른쪽부터 최소값을 기록하여 저장하는 배열 => O(N)
min_from_right = [0] * (n + 1)
current_min = 10_001  # 10001과 동일
for i in range(n, -1, -1):
    min_from_right[i] = min(current_min, arr[i])
    current_min = min_from_right[i]

# 전체 합과 최대 평균 계산
total_sum = sum(arr)
max_avg = 0

# 왼쪽부터 k개 원소를 삭제하면서 최대 평균 찾기 => O(N)
for k in range(1, n - 1):
    remaining_sum = total_sum - prefix_sum[k] - min_from_right[k + 1]
    remaining_avg = remaining_sum / (n - k - 1)
    max_avg = max(max_avg, remaining_avg)

# 결과 출력
print(f"{max_avg:.2f}")