나무 심기

백준 1280번: 나무 심기

아이디어

해당 구간에 존재하는 나무의 개수와 인덱스 0부터 각 나무까지 거리의 합을 세그먼트 트리에 기록한다.
이렇게 하면 새로운 나무를 심을 때 바로 cost를 구할 수 있다.

리프 노드에는 {인덱스 번호 * 나무 개수, 나무 개수}가 기록되어 있다.
새로 심는 나무의 인덱스 i를 기준으로
왼쪽은 left cnt * i - left sum, 오른쪽은 right sum - right cnt * i가 전체 비용이 된다.

코드

#include <bits/stdc++.h>

using namespace std;

const int MAX = 200000;
const int MOD = 1000000007;
int N;

vector<pair<long, long>> tree(MAX*4, {0, 0}); // sum, count

pair<long, long> query(int start, int end, int left, int right, int node) {
    if (right < start || end < left) return {0, 0};
    else if (left <= start && end <= right) return tree[node];
    pair<long, long> p1 = query(start, (start+end)/2, left, right, node*2);
    pair<long, long> p2 = query((start+end)/2+1, end, left, right, node*2+1);
    return {p1.first+p2.first, p1.second+p2.second};
}

void update(int start, int end, int idx, int node) {
    if (idx < start || end < idx) return;
    if (start == end) {
        tree[node].first += idx;
        tree[node].second++;
        return;
    }
    update(start, (start+end)/2, idx, node*2);
    update((start+end)/2+1, end, idx, node*2+1);
    tree[node].first = tree[node*2].first + tree[node*2+1].first;
    tree[node].second = tree[node*2].second + tree[node*2+1].second;
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> N;
    long long ans = 1;
    for (int i = 0; i < N; i++) {
        pair<long, long> ret;
        int x;
        long long tmp = 0;
        cin >> x;

        if (i == 0) {
            update(0, MAX-1, x, 1);
            continue;
        }

        ret = query(0, MAX-1, 0, x-1, 1);
        tmp += ret.second*x - ret.first;
        ret = query(0, MAX-1, x+1, MAX-1, 1);
        tmp += ret.first - ret.second*x;

        ans %= MOD;
        ans *= tmp%MOD;
        update(0, MAX-1, x, 1);
    }
   
    cout << ans%MOD;
    return 0;
}

여담

정답률이 16%다. 다들 오버플로우가 나서 틀린듯. 거리 합은 long long으로 선언하고 답은 계속 MOD로 잘 나눠주자. 다 풀고 다른사람들 짠거 보니까 다들 펜윅 트리로 풀었던데 나도 저거 공부해야겠다.