BOJ 백준 16296번 Daily division

BAPC 2018 Preliminary D번 문제이다.

풀이에 앞서 문제를 먼저 정리하자면, 1번 부터 N번까지의 마을이 일렬로 있으며 각 마을에 몇 명이 있는지가 주어진다. 날마다 푸드트럭이 이 마을로 찾아오는데 i번 마을에 찾아왔을 때 1번부터 i-1번 마을까지의 사람들이 왼쪽에 줄을 서고, i+1번부터 N번 마을까지의 사람들이 오른쪽에 줄을 서게 된다. i번 마을에 있는 사람들은 반으로 나눠져 각각 왼쪽 오른쪽에 서며, 사람 수가 홀수일 경우에는 남는 1명은 양쪽 줄의 차이가 적어지는 쪽으로 이동하되, 양쪽 사람 수가 같으면 아무 곳이나 간다.

이 때, 쿼리는 i번 마을의 사람 수가 임의로 바뀌었을 때 양쪽 줄의 차이가 가장 적어지는 푸드 트럭의 위치는 어디인가이다.

트럭이 i번 마을에 왔을 때, 왼쪽 줄 사람수에서 오른쪽 줄 사람수를 뺀 것을 $ F(i)$라고 하고, i번 마을의 사람 수를 $a_i$라고 하자 그러면,

$$
F(i) = \Sigma_{j=1}^{i-1}{a_j} - \Sigma_{j=i+1}^{N}{a_j}
$$

로 나오고 i에 대한 단조 증가 함수임을 쉽게 알 수 있다. 또한, $F(1) < 0$이며, $F(N) > 0$임은 자명하다. 왼쪽 줄과 오른쪽 줄의 사람 수 차이는 $\lvert F(i) \rvert$이므로 $F(i) \ge 0$이면서 가장 작은 $i$를 찾으면 $i$ 혹은 $i-1$이 답이 됨을 알 수 있다.

그러한 $i$를 찾는 것은 이분탐색을 통해 가능하며, 수행시간은 $O(logN)$ 이다.

이제 위에서 배제했던 i번 마을의 사람들을 고려하자. i번 마을의 사람 수가 홀수일경우 양쪽의 차이를 줄이는 방향으로 움직인다고 했다. 그리고 같을 경우에는 임의로 선택한다고 했으니 수가 같을 때는 항상 왼쪽으로 간다고 하자. 그럼 위 $F$식이 아래와 같이 변한다.

$$
F(i) = \Sigma_{j=1}^{i-1}{a_j} - \Sigma_{j=i+1}^{N}{a_j} + \alpha \text{ where } \alpha={1\text{ if } \Sigma_{j=1}^{i-1}{a_j} \leq \Sigma_{j=i+1}^{N}{a_j}\brace -1 \text{ if } \Sigma_{j=1}^{i-1}{a_j} \gt \Sigma_{j=i+1}^{N}{a_j}}
$$

물론 식이 이렇게 변하여도 답은 위에서 말했듯이 $F(i)\ge0$이면서 가장 작은 $i$를 찾았을 때, $i$ 혹은 $i-1$인 것은 변하지 않는다.

위 식을 실제로 구하기 위해서는 원소가 바뀔 때에도 구간 합을 빠르게 구할 수 있어야 한다. 이를 위해서는 세그먼트 트리를 이용하면 각 구간합을 $O(logN)$만에 구할 수 있다.

$F(i)\ge0$인 $i$를 구하는데에 걸리는 시간은 이분탐색에 $O(logN)$, $F$값을 구하는 데에 $O(logN)$으로 쿼리 하나당 $O(log^2N)$이 걸리게 된다. 따라서 총 시간복잡도는 $O(NlogN + Qlog^2N)$이 된다.

아래는 코드다. 다른 거 없고 구간 합만 구하면 되기 때문에 코드가 나한테 좀 더 익숙한 펜윅트리를 사용했다.

#include<bits/stdc++.h>
using namespace std;
using ll = long long;

ll arr[100005];
int N, Q;

ll query(vector<ll> &tree, int idx) {
    ll ret = 0;
    while(idx > 0) {
        ret += tree[idx];
        idx -= (idx & -idx);
    }
    return ret;
}

void update(vector<ll> &tree, int idx, ll diff) {
    while(idx < tree.size()) {
        tree[idx] += diff;
        idx += (idx & -idx);
    }
}

ll total;

ll f(vector<ll> &tree, ll i) {
    ll left = query(tree, i-1);
    ll right = total - left - arr[i];
    if(arr[i] % 2) {
        if(left <= right) ++left;
        else ++right;
    }
    return left - right;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> Q;
    vector<ll> tree(N+1);
    for(int i=1;i<=N;++i) {
        cin >> arr[i];
        total += arr[i];
        update(tree, i, arr[i]);
    }
    for(int i=0;i<Q;++i) {
        int a;
        ll x;
        cin >> a >> x;
        update(tree, a+1, x - arr[a+1]); total += x - arr[a+1];
        arr[a+1] = x;
        int lo = 1, hi = N;
        while(lo < hi) {
            int mid = (lo+hi) / 2;
            ll val = f(tree, mid);
            if(val < 0) lo = mid+1;
            else hi = mid;
        }
        if(lo > 1 && -f(tree, lo-1) <= f(tree, lo)) --lo;
        cout << lo-1 << '\n';
    }
    return 0;
}