프로그래머스 지형 편집 (C++)

프로그래머스

1*1*1크기의 정육면체 블록으로 만들어진 지형의 상태를 입력 받습니다. 정육면체 하나를 추가하는 비용 P, 제거하는 비용이 Q일 때, 지형을 평평하게 만드는 최소 비용이 얼마인지 구해주세요.

 

 

풀이 방법

기준이 되는 높이를 찾고 해당 높이와 다른 곳에 공사를 하여 평평하게 맞춰주는 비용을 반환하면 됩니다.

여기서 문제가 되는 부분은, 기준이 되는 높이를 찾는 것과 높이가 다른 곳의 비용을 모두 더하는 것 입니다.

이 문제를 해결하지 않으면 시간초과가 발생합니다.

 

 

높이가 다른 곳의 비용 구하기 - 누적합

이 부분을 빠르게 해결하기 위해서는 누적합을 이용해야 합니다.

for (auto& v : land)
    for (int& n : v)
        heights.push_back(n);

sort(heights.begin(), heights.end());
prefixSum.push_back(0);

for (int i = 0; i < heights.size(); i++)
    prefixSum.push_back(prefixSum[i] + heights[i]);

2차원 배열로 입력받은 land를 1차원 배열로 옮긴 후, 추가적인 배열을 만들어서 누적합을 구해줍니다. 

이러한 높이의 누적합을 이용하면 기준이 되는 높이보다 작은 높이의 블록 수를 한번에 구할 수 있게 되고,

더 높은 블록의 수를 한번에 구할 수 있게 됩니다.

 

long long GetCost(ll h)
{
    int idx = lower_bound(heights.begin(), heights.end(), h) - heights.begin();
    
    ll lowerSum = prefixSum[idx];
    ll upperSum = prefixSum.back() - lowerSum;
    
    ll lowerCnt = h * idx - lowerSum;
    ll upperCnt = upperSum - h * (heights.size()-idx);
    
    return lowerCnt * P + upperCnt * Q;
}

1. lower_bound를 통해 정렬돼있는 heights 배열에서 h의 위치값을 찾아냅니다.

2. 더 낮은 인덱스의 블록 수 는 prefixSum[idx]를 통해 구할 수 있습니다.

3. 더 높은 인덱스의 블록 수는 모든 블록의 수에 lowerSum을 빼면 구할 수 있습니다.

4. 쌓아올리는 블록의 수는 목표가 되는 블록 수 (h*idx)에 lowerSum을 빼면 구할 수 있고

5. 제거할 블록의 수는 upperSum에  (h* (heights.size()-idx))를 빼서 구할 수 있게 됩니다.

6. 이후 쌓는 비용과 제거 비용을 곱해서 반환하면 logN의 속도로 해결 됩니다 (lower_bound)

 

 

기준이 되는 높이 빠르게 찾기 - 삼분 탐색

이분 탐색이 아닌, 삼분 탐색을 이용하여 값을 찾는 이유는 비용이 U자 함수를 그리기 때문입니다.

해당 문제는 높이가 높아질 수록 비용이 한 방향으로 바뀌는게 아닌, 양 방향에서 한쪽으로 줄어드는 U자형 구조입니다.

 

그렇기 때문에 정렬이 된 수열 같은 곳에서 활용도가 좋은 이분탐색을 해당 문제에 적용하기는 힘듭니다.

 

long long TernarySearch(ll left, ll right)
{
    while (right - left > 2)
    {
        ll mid1 = left + (right - left) / 3;
        ll mid2 = right - (right - left) / 3;
        
        ll cost1 = GetCost(mid1);
        ll cost2 = GetCost(mid2);
        
        if (cost1 < cost2)
            right = mid2;
        else
            left = mid1;
    
    }
    
    ll minCost = LLONG_MAX;
    for (ll h = left; h <= right; h++)
        minCost = min(minCost, GetCost(h));
    
    return minCost;
}

left와 right를 기준으로 2개의 중간 인덱스를 이용하여 높이별 완성 비용을 구합니다.

이후, 비용을 비교하여 범위를 좁혀가는 방식입니다.

이후 남아있는 영역(left - right)의 비용을 구하여 최소 비용을 찾아주면 됩니다.

 

 

 

 

정답 코드

#include<vector>
#include<algorithm>
#include<iostream>
#include <climits>

using namespace std;
typedef long long ll;

vector<ll> heights;
vector<ll> prefixSum;
ll P, Q;

long long GetCost(ll h)
{
    int idx = lower_bound(heights.begin(), heights.end(), h) - heights.begin();
    
    ll lowerSum = prefixSum[idx];
    ll upperSum = prefixSum.back() - lowerSum;
    
    ll lowerCnt = h * idx - lowerSum;
    ll upperCnt = upperSum - h * (heights.size()-idx);
    
    return lowerCnt * P + upperCnt * Q;
}

long long TernarySearch(ll left, ll right)
{
    while (right - left > 2)
    {
        ll mid1 = left + (right - left) / 3;
        ll mid2 = right - (right - left) / 3;
        
        ll cost1 = GetCost(mid1);
        ll cost2 = GetCost(mid2);
        
        if (cost1 < cost2)
            right = mid2;
        else
            left = mid1;
    
    }
    
    ll minCost = LLONG_MAX;
    for (ll h = left; h <= right; h++)
        minCost = min(minCost, GetCost(h));
    
    return minCost;
}

long long solution(vector<vector<int> > land, int p, int q) {
    P = p; Q = q;
    
    for (auto& v : land)
        for (int& n : v)
            heights.push_back(n);
    
    sort(heights.begin(), heights.end());
    prefixSum.push_back(0);
    
    for (int i = 0; i < heights.size(); i++)
        prefixSum.push_back(prefixSum[i] + heights[i]);
    
    return TernarySearch(heights[0], heights.back());
}