Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 여러 개의 막대를 하나로 연결하는 최소 비용 구하기

길이가 양의 정수인 막대(stick) 여러 개가 주어져 있다고 가정해 봅시다. 길이가 각각 X와 Y인 두 막대를 하나로 이을 때 비용은 X + Y가 되며, 이 과정은 막대가 하나만 남을 때까지 반복됩니다. 우리가 구해야 할 것은 모든 막대를 하나로 연결할 때 드는 최소 비용입니다.

예를 들어 막대 배열이 [2, 4, 3]이라면, 정답은 14가 됩니다.

해결 전략: 그리디 알고리즘과 최소 힙(Min Heap)

연결 비용이 항상 두 막대 길이의 합이기 때문에, 매 단계에서 가장 짧은 막대 두 개를 먼저 연결하는 것이 전체 비용을 최소화하는 핵심 아이디어입니다. 이러한 선택을 효율적으로 하기 위해 최소 힙(min heap) 기반 우선순위 큐(priority queue)를 사용합니다.

알고리즘 단계

  • 최소 힙으로 동작하는 우선순위 큐 pq를 선언합니다.
  • 배열 s의 모든 원소를 pq에 삽입합니다.
  • 누적 비용 변수 ans를 0으로 초기화합니다.
  • pq에 원소가 두 개 이상 남아 있는 동안 다음을 반복합니다.
    • 큐의 최상단 값을 temp에 저장하고 제거합니다.
    • temp에 새로운 최상단 값을 더한 뒤 해당 값도 제거합니다.
    • ans에 temp를 더해 누적 비용을 갱신합니다.
    • temp(새로 만들어진 막대의 길이)를 다시 pq에 삽입합니다.
  • 반복이 끝나면 ans를 반환합니다.

C++ 예제 코드

다음 구현 예제를 통해 동작 방식을 더 자세히 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
    int connectSticks(vector<int>& s) {
        // greater<int>를 사용해 최소 힙(min heap)으로 구성
        priority_queue <int, vector<int>, greater<int> > pq;
        for(int i = 0; i < s.size(); i++) pq.push(s[i]);
        int ans = 0;
        while(pq.size() > 1){
            int temp = pq.top();
            pq.pop();
            temp += pq.top();
            pq.pop();
            ans += temp;
            pq.push(temp);
        }
        return ans;
    }
};
main(){
    vector<int> v = {2, 4, 3};
    Solution ob;
    cout << ob.connectSticks(v);
}

입력

[2,4,3]

출력

14

동작 과정 살펴보기

[2, 4, 3] 입력값이 처리되는 과정을 단계별로 확인하면 다음과 같습니다.

  1. 가장 짧은 두 막대 2와 3을 연결 → 비용 5 발생, 누적 비용 = 5. 새 막대의 길이는 5.
  2. 남은 막대 중 가장 짧은 두 개인 4와 5를 연결 → 비용 9 발생, 누적 비용 = 14.

모든 막대가 하나로 합쳐졌으므로 최종 결과는 14입니다. 매번 가장 작은 값부터 연결하면 긴 막대가 여러 번 더해지는 것을 방지할 수 있어 총비용이 최소화됩니다.

시간 복잡도

막대가 n개일 때, 힙 연산(push/pop) 한 번당 O(log n)의 시간이 걸리고 이를 총 n-1회 수행하므로 시간 복잡도는 O(n log n)입니다. 공간 복잡도는 힙에 모든 막대를 저장해야 하므로 O(n)입니다.