[코드포스 Practice19] D. Minimize the error

아 이거 열심히 생각했는데 입력 1000밖에 안 돼서 간단한 방법으로 풀 수 있었다.

이게 왜 D번이지??

 

#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <functional>
#include <string>
#include <queue>
#include <stack>
#include <set>
 
using namespace std;
using i64 = long long;
using ii = pair<int, int>;
using ii64 = pair<i64, i64>;
 
int main() 
{
    int n, k1, k2;
    scanf("%d %d %d", &n, &k1, &k2);
 
    vector<int> a(n);
    vector<int> b(n);
    for (int i = 0; i < n; i++)
        scanf("%d", &a[i]);
    for (int i = 0; i < n; i++)
        scanf("%d", &b[i]);
 
    vector<int> sub(n);
    for (int i = 0; i < n; i++)
        sub[i] = abs(a[i] - b[i]);
    
    sort(sub.begin(), sub.end(), greater<int>());
    
    // for (int i = 0; i < n; i++)
    //     cout << sub[i] << " ";
    // cout << endl;
    
    i64 sum = k1 + k2;
    // cout << sum << endl;
    for (int i = 0; i < n; i++)
    {
        if (sub[i] == 0)
            break;
        if (sum - sub[i] <= 0)
        {
            sub[i] -= sum;
            sum = 0;
            break;
        }
        else
        {
            sum -= sub[i];
            sub[i] = 0;
        }
    }
    
    // cout << sum << endl;
    // for (int i = 0; i < n; i++)
    //     cout << sub[i] << " ";
    
    if (sum != 0)
    {
        // cout << sum << " " << n << endl;
        cout << (sum%n)*(sum/n+1)*(sum/n+1) + (n - sum%n)*(sum/n)*(sum/n);
    }
    else
    {
        i64 ans = 0;
        for (int i = 0; i < n; i++)
            ans += sub[i] * sub[i];
        cout << ans;
    }
    
    return 0;
}

기존 풀었던 코드는 이렇고

5 5 5

0 0 0 0 0

0 0 0 0 0

여기서 막혔다

 

#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <functional>
#include <string>
#include <queue>
#include <stack>
#include <set>
 
using namespace std;
using i64 = long long;
using ii = pair<int, int>;
using ii64 = pair<i64, i64>;
 
int main() 
{
    int n, k1, k2;
    scanf("%d %d %d", &n, &k1, &k2);
 
    vector<int> a(n);
    vector<int> b(n);
    for (int i = 0; i < n; i++)
        scanf("%d", &a[i]);
    for (int i = 0; i < n; i++)
        scanf("%d", &b[i]);
 
    vector<i64> sub(n);
    for (int i = 0; i < n; i++)
        sub[i] = abs(a[i] - b[i]);
    
    i64 sum = k1 + k2;
    for (int i = sum; i > 0; i--)
    {
        int max_idx = 0;
        for (int i = 0; i < n; i++)
        {
            if (sub[i] > sub[max_idx])
                max_idx = i;
        }
        if (sub[max_idx] > 0)
            sub[max_idx]--;
        else
            sub[max_idx]++;
    }
    
    i64 ans = 0;
    for (int i = 0; i < n; i++)
    ans += sub[i] * sub[i];
    cout << ans;
    
    return 0;
}

 

으음.., 정직하게 계속 배열 돌면서 1씩 더하고 빼고 했다.

 

이거 O(N^2)인데 최대값 구하는 부분을 logN으로 줄이면 O(NlogN)까지 가능하다.

 

priority queue쓰면 logN에 가능하다

 

#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <functional>
#include <string>
#include <queue>
#include <stack>
#include <set>
 
using namespace std;
using i64 = long long;
using ii = pair<int, int>;
using ii64 = pair<i64, i64>;
 
int main() 
{
    int n, k1, k2;
    scanf("%d %d %d", &n, &k1, &k2);

    vector<int> a(n);
    vector<int> b(n);
    for (int i = 0; i < n; i++)
        scanf("%d", &a[i]);
    for (int i = 0; i < n; i++)
        scanf("%d", &b[i]);

    priority_queue<i64> sub;
    for (int i = 0; i < n; i++)
        sub.push(abs(a[i] - b[i]));
    
    i64 sum = k1 + k2;
    for (int i = sum; i > 0; i--)
    {
        int top = sub.top();
        sub.pop();
        if (top > 0)
        {
            sub.push(top-1);
        }
        else
        {
            sub.push(top+1);
        }
    }
    
    i64 ans = 0;
    for (int i = 0; i < n; i++)
    {
        ans += sub.top() * sub.top();
        sub.pop();
    }
    cout << ans;
    
    return 0;
}

우선순위 큐 써봤다

예아