백준 21916번 Neo-Robin Hood

문제 요약

돈을 $m_i$달러만큼 들고 있는 정치가 $N$명이 주어진다. 각 정치가는 $p_i$만큼 돈을 원하고 있다. 이 때, 정치가들한테서 돈을 훔치려고 하는데 그냥 쌩으로 훔치면 위험하기 때문에 훔친사람의 수와 같은 수의 정치가들한테 그들이 원하는 만큼의 돈을 줘야 한다.

이 때, 돈을 훔칠 수 있는 정치가의 수의 최댓값을 구하시오.

$ N \le 100,000$ $1 \le m_i, p_i \le 10^9$

풀이

관찰 1 : 일단 돈을 줄 사람 $k$ 명을 정했다고 하자. 그렇게 되면 돈을 훔쳐야 될 사람 $k$ 명은 남은 사람 중에서 돈이 제일 많은 $k$ 명이 되는 것은 자명하다.

관찰 2. $k$ 명에게서 훔치는 것이 가능하다면 당연히 $k$ 명보다 적은 수의 정치가에게서 훔치는 것도 가능하다. 이를 통해 이분탐색이 가능하다.

$k$가 고정되었을 때 가능한지의 결정문제를 적절히 빠르게 풀면 문제가 해결 가능해보인다.

이제 관찰 1을 통해서 $k$ 명에게서 돈을 훔치려고 할 때 그게 가능한지에 대해 판별하는 결정문제를 해결하자.

먼저 정치가들을 $m_i$의 내림차순으로 정렬하자. 그리고 돈을 훔칠 사람을 1부터 i번째 사람중에서 k명 뽑는다고 생각을 해보자.

그러면 관찰 1을 통해서 i명 중에서 k명에게선 돈을 훔치고 i-k명에겐 돈을 줘야 되는 것을 알 수 있다. 그럼 $i+1$ 번째 사람부터 $N$ 번째 사람중에선 돈을 줄 사람 $2k-i$ 명을 골라야 한다. 이 때 어떻게 사람을 골라야 최적일까?

먼저 알기 쉬운 $[i+1, N]$ 에 속하는 사람들부터 살펴보자. 전부 돈을 줄 사람들이기 때문에 그냥 $p_i$ 의 합이 가장 작게 되는 $2k-i$ 명을 고르면 된다. 이 때 이 합은 세그트리를 쓰면 $O(logN)$에 구할 수 있다.

이제 문제는 $[1, i]$ 에 속하는 사람들이다. 어떻게 골라야 비용을 최대한 많이 남길 수 있을까? $i$ 명 중에서 우리가 돈을 훔칠 사람은 $m_i$ 만큼의 이득을 주고 돈을 줄 사람은 $p_i$ 만큼의 손해를 준다. 이를 반대로 처음에 전부 손해라고 생각해보자.

그렇게 되면 $m_i$ 만큼 돈을 가진 정치가에게서 돈을 훔칠 때 $m_i + p_i$ 의 이득이 생기는 것과 동일하다. 따라서 $[1,i]$에서 그냥 $m_i+p_i$ 의 합이 가장 큰 $k$ 명을 고르면 그것이 최적이 된다. 이거는 그냥 우선순위 큐로도 가능하다.

이제 $[1, i]$ 에서 $k$ 명의 사람들에게서 돈을 훔친다고 했을 때 최적으로 돈을 뺏는 방법을 알았다. 그리고 그 때의 비용을 구할 수 있다.

이걸 $k \le i \le 2k$ 에 대해서 시도해보고 남은 돈이 0 이상인 적이 한 번이라도 있다면 $k$ 명에게서 돈을 뺏는 것이 가능한것을 알 수 있고, 그런 적이 없었다면 불가능한 것이다.

이걸로 결정문제를 $O(NlogN)$ 에 풀 수 있다. 이분탐색까지 하면 $O(logN)$이 곱해져서 총 시간복잡도 $O(Nlog^2N)$에 해결이 가능하다.

백준 문제풀이를 거의 한달간 안썼는데 팀연습을 한달만에 다시 했더니 바로 풀이를 쓸 게 생겼다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
int N;
vector<pair<ll,ll>> a;
ll psum[100005];
vector<ll> v;
pair<int, ll> tree[1<<18];

void update(int node, int s, int e, int idx, ll val) {
    if(e < idx || s > idx) return ;
    if(idx == s && e == idx) {
        tree[node].first += (val < 0? -1 : 1);
        tree[node].second += val;
        return ;
    }
    int mid = s + e >> 1;
    int n1 = node << 1, n2 = node << 1 | 1;
    update(n1, s, mid, idx, val);
    update(n2, mid+1, e, idx, val);
    tree[node].first = tree[n1].first + tree[n2].first;
    tree[node].second = tree[n1].second + tree[n2].second;
}

ll query(int k) {
    ll ret = 0;
    int node = 1;
    int s = 0, e = v.size()-1;
    while(s != e) {
        int mid = s + e >> 1;
        int n1 = node << 1, n2 = node << 1 | 1;
        int lsz = tree[n1].first;
        if(lsz >= k) {
            e = mid;
            node = n1;
        }
        else {
            ret += tree[n1].second;
            k -= lsz;
            s = mid + 1;
            node = n2;
        }
    }
    ret += v[s] * min(k, tree[node].first);
    return ret;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N;
    a.resize(N+1);
    a[0] = {1e9, 1e9};
    // first : m, second : p
    for(int i=1;i<=N;++i) cin >> a[i].first;
    for(int i=1;i<=N;++i) cin >> a[i].second;
    sort(a.begin(), a.end(), greater<>());
    for(int i=1;i<=N;++i) {
        psum[i] = psum[i-1] + a[i].second;
        v.push_back(a[i].second);
    }
    sort(v.begin(), v.end());
    v.erase(unique(v.begin(), v.end()), v.end());
    vector<int> v_idx(N+1);
    for(int i=1;i<=N;++i) {
        v_idx[i] = lower_bound(v.begin(), v.end(), a[i].second) - v.begin();
    }
    int lo = 0, hi = N/2 + 1;
    while(lo + 1 < hi) {
        int k = (lo + hi) >> 1;
        priority_queue<ll, vector<ll>, greater<>> pq;
        memset(tree, 0, sizeof(tree));
        ll lsum = 0;
        for(int i=1;i<k;++i) {
            pq.push(a[i].first + a[i].second);
            lsum += a[i].first + a[i].second;
        }
        for(int i=k;i<=N;++i) update(1, 0, v.size()-1, v_idx[i], a[i].second);
        bool flag = false;
        for(int i=k;i<=2*k;++i) {
            {
                ll x = a[i].first + a[i].second;
                if(pq.size() < k) {
                    pq.push(x);
                    lsum += x;
                }
                else if(pq.top() < x) {
                    lsum += x - pq.top();
                    pq.pop();
                    pq.push(x);
                }
            }
            ll rsum;
            {
                update(1, 0, v.size()-1, v_idx[i], -a[i].second);
                rsum = query(2*k-i);
            }
            if(lsum - rsum - psum[i] >= 0) {
                flag = true;
                break;
            }
        }
        if(flag) lo = k;
        else hi = k;
    }
    cout << lo << '\n';
    return 0;
}