[BOJ-1280]나무심기

이 문제 풀기까지 5일정도 걸린 것 같다. 어떻게 값을 구하는 방법은 알았으나.. "이걸 어떻게 펜윅트리로 만들어?" 라는 생각으로 계속 생각만 했던 문제였다. 5번 정도 틀렸는데 펜윅트리 때문이 아니라 나무가 가진 cost 가 long long 타입 값 범위를 넘어설 거라고 생각안하고 코드를 작성한 이유였다... 

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

[문제요약]

나무를 심는다.

나무를 심을 때 비용이 발생한다.

[조건 1] 단, 첫번째 나무는 비용이 발생하지 않는다. 

[조건 2] 1번째 나무를 제외한 i ( i >= 2) 번째 나무는 이전에 심어진 [1, i-1] 번째 나무들과 거리들의 합이 비용이 된다. 

그리고 1번째 심는 나무를 제외한 나머지 나무들 [2, n] 의 비용의 곱을 1,000,000,007 로 나눈 값을 출력으로 내놓으면 된다.

[입력]

n ~ [2, 200,000]

a1

a2

...

ai

...

an

첫번째 입력으로 나무의 개수

두번째부터 n+1번째까지 1번째 나무, ... , n번째 나무를 심는 위치가 입력으로 주어진다.

[출력]

1번째를 제외한 나머지 나무들의 비용의 곱을 1,000,000,007 로 나눈 값을 출력한다. 

[접근방법]

문제를 보고서, 첫날에 이것을 어떻게 하면 빨리 구할 수 있을까? 라는 생각부터 했었다. 

무조건 O(N^2)로 풀게되면 입력이 200,000 이므로 노답이다.

그래서 알고리즘 분류를 봤다. 알고리즘 분류는 '펜윅트리' 인 것을 보고 펜윅트리를 어떤 형식으로 사용해야 할까.. 했는데 생각이 나지 않았다.

그래서 각 나무의 비용을 구하는 식을 먼저 작성해보았다.

각 나무들의 비용: s1 s2 s3 ... si ... sn-1 sn 

각 나무들의 위치: a1 a2 a3 ... ai ... an-1 an

[조건 1]에 따라 s1은 곱하지 않아도 된다.

그럼 s2 부터 차례대로 구해보자.

s2 = |a2 - a1|

s3 = |a3 - a1| + |a3 - a2|

s4 = |a4 - a1| + |a4 - a2| + |a4 - a3|

....

si = |ai - a1 | + |ai - a2| + ... + |ai - ai-1|

....

sn = |an - a1| + ... + |an - an-1|

으로 각 s 값들을 구할 수 있다.

그럼 si 의 값은 아래 식과 같아진다.

 

si =

 

이때 |ai - ak| 는 3가지 경우로 나눌 수 있다. 

1. 현재 i번째 심는 나무 위치보다 k번째 심은 나무 위치가 큰 경우 ~ if ( ai < ak ) ~ i보다 더 큰 위치에 있는 경우
2. 현재 i번째 심는 나무 위치보다 k번째 심은 나무 위치가 작은 경우 ~ if ( ai > ak ) ~ i보다 더 작은 위치에 있는 경우
3. 현재 i번째 심는 나무 위치와 k번째 심은 나무 위치가 같은 경우 ~ if ( ai == ak )

          이때 3번의 경우는 0이므로 고려하지 않음.

1번의 경우에는 |ai - ak| == ak - ai 와 같다. 

2번의 경우에는 |ai - ak| == ai - ak 와 같다. 

위 경우를 보았을 때 작은 것과 큰 것을 구분해서 놓을 수 있다. 

si = = ( a큰경우위치의 합 - 큰경우 * ai ) + (작은경우 * ai - a작은위치의 합) + (같은 경우* ai - ai 같은경우) 

이때 같은 경우 값은 0이므로 제외시킨다. 

그러면 큰 경우의 합과 개수, 작은 경우의 합과 개수를 구하면 i번째 심는 나무의 비용을 알 수 있다. 

여기까지가 문제를 보고 30분동안 생각한 내용이었다. 

펜윅을 어떻게 적용시켜야 할 지가 문제였다. 

... ... ... 그리고 4일이 지났다.

구해야하는 것은 

1. 각 경우의 합
2. 각 경우의 개수

그러면 두 개의 펜윅트리를 만들어야 하나? 라고 생각했다.

합에 대한 펜윅트리, 개수에 대한 펜윅트리를 만들면 된다. 

단, 펜윅트리는 인덱스 [0] 을 사용할 수 없다. 

그러기 때문에 각 자리의 합을 구할 때 한 칸(+1) 씩 밀어줘야 한다는 걸 알았다.

대신 합에 대한 펜윅트리를 만들 때는 원래 값을 넣어줘야 하기때문에 (-1)을 한 값을 넣어준다. 

count 0 1 2 3 4 5 6 7 8 9 .... 16 .... 200000 ....

sum 0 1 2 3 4 5 6 7 8 9 .... 16 .... 200000 ....

각 인덱스 위치는 나무를 심는 위치를 의미하게 된다. 

만약 0번 위치에 나무를 심는 경우 

count 는 +1 를 하는데 0번 인덱스를 사용할 수 없으므로, +1 한 1번 인덱스부터 사용할 수 있다. 

합에 대한 펜윅트리는 인덱스 사용은 동일하다. 하지만 더해지는 값은 나무를 심는 위치값이 들어가게 된다.

위 내용대로 코드를 작성하였다.

#include<cstdio>
#include<vector>
#define din1(a) scanf("%d", &a)
#define llin1(a) scanf("%lld", &a)
#define pb(a) push_back(a)
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
 
const int MMAX = 440000;
const int INF = 202020;
const ll DIV = 1000000007;
class BIT {
public:
    vll T;
    vll A;
 
    BIT() : T(MMAX), A(MMAX, 0) {}
 
    void update(int p, ll v) {
        while (p <= INF) {
            T[p] += v;
            p += p& (-p);
        }
    }
    
    ll sum(int p) {
        ll ret = 0;
        while (p > 0) {
            ret += T[p];
            p -= p & (-p);
        }
        return ret;
    }
};
vll place;
int main() {
    int n;  din1(n);
    BIT count_fw = BIT();
    BIT sum_fw = BIT();
    place = vll(MMAX, 0);
    ll ans = 1;
    for (int i = 1; i <= n; i++) {
        llin1(place[i]);
        place[i]++;
 
        count_fw.update(place[i], 1);
        sum_fw.update(place[i], place[i] - 1);
        if (i == 1) continue; 
        ll upper = count_fw.sum(INF) - count_fw.sum(place[i]-1);
        ll lower = count_fw.sum(place[i]-1);
 
        ll S_upper = sum_fw.sum(INF) - sum_fw.sum(place[i]-1);
        ll S_lower = sum_fw.sum(place[i]-1);
        
        ll pivot = place[i] - 1;
        ll cost = (S_upper - pivot*upper + pivot*lower - S_lower) % DIV; // 값이 ll 을 넘칠 수 있음.. 주의...
        ans = (ans * cost) % DIV;
        //printf("pivot : %lld , answer : %lld / %d upper : %lld (%lld), lower : %lld (%lld)\n",pivot,ans, i, upper, S_upper, lower, S_lower);
    }
    if (n == 1) ans = 0;
 
    printf("%lld\n", ans);
    return 0;
}