Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 K명의 직원을 최소 비용으로 고용하는 프로그램 작성하기

문제 설명

각 직원의 능력을 담은 배열 quality와 각 직원의 최저 임금 기대치를 담은 배열 wage, 그리고 정수 K가 주어집니다. i번째 직원은 quality[i]만큼의 능력을 가지고 있으며, 최소한 wage[i] 이상의 임금을 받기를 기대합니다. 우리는 K명의 직원을 고려하여 하나의 유급 그룹을 구성하려고 하며, 이때 다음 두 가지 규칙을 반드시 지켜야 합니다.

  • 유급 그룹에 속한 모든 직원은 그룹 내 다른 직원들과 비교했을 때 자신의 능력(quality) 비율에 맞게 급여를 받아야 합니다.
  • 유급 그룹에 속한 모든 직원은 최소한 자신의 최저 임금 기대치(wage) 이상을 받아야 합니다.

목표는 위 조건을 모두 만족하면서 K명의 직원으로 유급 그룹을 구성할 때 드는 최소 비용을 구하는 것입니다.

예를 들어 quality = [10, 22, 5], wage = [70, 52, 30], K = 2라고 가정해 보겠습니다. 이 경우 첫 번째 직원에게 70을, 세 번째 직원에게 35를 지급하면 조건을 충족하므로 정답은 105.000이 됩니다.

풀이 접근 방법

이 문제는 각 직원의 '임금 대비 능력 비율'을 기준으로 정렬한 뒤, 힙(heap) 자료구조를 활용하면 효율적으로 해결할 수 있습니다. 단계별 절차는 다음과 같습니다.

  1. 새로운 리스트 qr을 생성합니다.
  2. quality와 wage의 각 쌍 (q, w)에 대해 비율 w/q와 능력 q를 묶어 qr의 끝에 추가합니다.
  3. 리스트 qr을 오름차순으로 정렬합니다.
  4. 힙으로 사용할 새로운 리스트 cand와 누적합 변수 csum := 0을 준비합니다.
  5. i를 0부터 K-1까지 반복하면서 -qr[i][1]을 힙 cand에 삽입하고, csum에 qr[i][1]을 더합니다.
  6. ans := csum * qr[K-1][0]으로 초기화합니다.
  7. idx를 K부터 quality의 크기까지 반복하면서 다음을 수행합니다.
    • -qr[idx][1]을 힙 cand에 삽입합니다.
    • csum에 qr[idx][1]을 더한 뒤, 힙에서 최상위 요소(저장된 값 중 가장 큰 능력의 음수)를 꺼내 함께 더합니다. 즉, 현재 그룹에서 가장 능력이 높은 직원을 제거하는 효과가 있습니다.
    • ans를 min(ans, csum * qr[idx][0])으로 갱신합니다.
  8. 최종적으로 ans를 반환합니다.

동작 원리

핵심 아이디어는 각 직원의 시급에 해당하는 값(wage/quality)을 계산해 오름차순으로 정렬하는 것입니다. 특정 직원을 기준으로 삼으면, 그 직원보다 비율이 낮거나 같은 직원들은 동일한 단가를 적용해도 최저 임금 조건을 자연스럽게 충족합니다. 따라서 각 기준점마다 K명을 선택할 때는 능력 합이 가장 작은 K명을 고르는 것이 최적이며, 이 선택을 최대 힙(음수 값을 저장하는 최소 힙)으로 관리하면 매번 가장 능력이 높은 직원을 O(log K)에 제거할 수 있습니다.

예제 코드

다음 구현 예시를 통해 더 잘 이해해 보겠습니다.

import heapq

def solve(quality, wage, K):
    qr = []
    for q, w in zip(quality, wage):
        qr.append([w/q, q])
    qr.sort()

    cand, csum = [], 0
    for i in range(K):
        heapq.heappush(cand, -qr[i][1])
        csum += qr[i][1]
    ans = csum * qr[K - 1][0]

    for idx in range(K, len(quality)):
        heapq.heappush(cand, -qr[idx][1])
        csum += qr[idx][1] + heapq.heappop(cand)
        ans = min(ans, csum * qr[idx][0])
    return ans

quality = [10, 22, 5]
wage = [70, 52, 30]
K = 2
print(solve(quality, wage, K))

입력

[10,22,5], [70,52,30], 2

출력

105