백준 1280 - 나무 심기

https://www.acmicpc.net/problem/1280

힌트 1

특정 구간에 존재하는 나무의 개수와 위치의 합을 보존한다.

 

특정 구간에 존재하는 나무의 개수를 활용해서 이전에 심어진 나무와의 거리 합을 빠르게 구하는 방법을 생각해본다.

힌트 2

i를 이전에 심은 나무의 index, X는 현재 나무를 심을 index라고 하자.

 

심을 나무와 심어진 나무들의 거리 합은 다음 식으로 표현할 수 있다.

 

$sum = \sum_{i = 1}^{X - 1}|tree_i - tree_X|$

 

해당 식은 다음과 같이 바꿀 수 있다.

 

$sum = |\sum_{i = 1}^{X - 1}tree_i - \sum_{i = 1}^{X - 1}tree_X|$

 

이 식을 구현할 방법을 생각해본다.

---

풀이

i를 index라 했을 때 각 나무 위치와의 거리 합 $sum = \sum_{i = 1}^{X - 1}|tree_i - tree_X|$ 을 빠르게 구하는 게 핵심이다.

 

수식을 빠르게 계산하기 위해 먼저 세그먼트 트리 두 개를 이용할 것이다.

 

하나는 구간에 나무가 몇 개 심어진 상태인지, 나머지는 구간에서 심어진 나무의 위치 합이다.

 

두 개의 세그먼트 트리를 이용하여 위를 수정한 다음 식을 빠르게 계산할 수 있다.

 

$sum = |\sum_{i = 1}^{X - 1}tree_i - \sum_{i = 1}^{X - 1}tree_X|$

 

우항의 두 번째 항 $\sum_{i = 1}^{X - 1}tree_X$ 을 보자. 이는 $tree_X$와 지금까지 심은 나무의 개수를 곱하면 간단히 구할 수 있다.

 

나무가 심어진 개수를 구하는 세그먼트 트리를 통해 구하자.

 

우항의 첫 번째 항 $\sum_{i = 1}^{X - 1}tree_i$ 은 위치의 합을 관리하는 세그먼트 트리를 통해 구할 수 있다.

 

역시 질의를 통해 구한 후 위에서 구한 결과를 빼서 수식의 결과를 얻을 수 있다.

 

이 과정을 나무를 심을 때마다 구한 값을 곱해주면 정답을 얻을 수 있다.

구현

다만 현재 심은 위치의 왼쪽, 현재 심은 위치의 오른쪽을 따로 계산해야 한다.

 

심을 나무의 왼쪽 구간은 나무의 개수 * 현재 위치 - 왼쪽 구간에서의 위치합으로 구할 수 있다.

 

오른쪽은 반대로 오른쪽 구간의 위치합 - 현재 위치 * 나무의 개수를 계산하면 된다.

 

이 둘을 합치면 현재 나무를 심는데 필요한 비용이 된다.

 

게산 결과를 오버플로우에 특히 유의하면서 곱해주도록 한다.

 

전체 코드

#pragma GCC target("avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>
#define mod 1000000007

using namespace std;
using ll = long long;

int n;

inline int mid(int s, int e) { return s + (e - s) / 2; }

void update(vector<ll> &tree, int node, int start, int end, int idx, int diff)
{
    if (idx < start || idx > end) return;
    tree[node] += diff;
    if (start != end)
    {
        int m = mid(start, end);
        update(tree, node * 2, start, m, idx, diff);
        update(tree, node * 2 + 1, m + 1, end, idx, diff);
    }
}

ll sum(vector<ll> &tree, int node, int start, int end, int l, int r)
{
    if (start > r || end < l) return 0;
    if (l <= start && end <= r) return tree[node];

    int m = mid(start, end);
    return sum(tree, node * 2, start, m, l, r) + sum(tree, node * 2 + 1, m + 1, end, l, r);
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin >> n;
    vector<ll> density(800000), dist(800000);

    ll res = 1;

    for (int i = 0; i < n; ++i)
    {
        int x;
        cin >> x;
        ++x;
        if (i)
        {
            ll l_cnt = 0, l = 0;
            if (x > 1)
            {
                l_cnt = sum(density, 1, 1, 200000, 1, x - 1);
                l = l_cnt * x - sum(dist, 1, 1, 200000, 1, x - 1);
                l %= mod;
            }
            
            ll r_cnt = 0, r = 0;

            if (x < 200000)
            {
                r_cnt = sum(density, 1, 1, 200000, x + 1, 200000);
                r = sum(dist, 1, 1, 200000, x + 1, 200000) - r_cnt * x;
                r %= mod;
            }
            
            res *= (l + r) % mod;
            res %= mod;
        }

        update(density, 1, 1, 200000, x, 1);
        update(dist, 1, 1, 200000, x, x);
    }

    cout << res;
}