[백준 1280번] [세그먼트 트리] 나무 심기

Approach

문제를 잘 보면, cost를 결정하기 위해서는 궁극적으로 나무를 심는 위치를 기반으로 앞에 있는 나무의 개수와 어디에 위치해 있는지 판단하는 것이 중요하다.

 

Let $i$번재 나무 위치를 $x_i$라 하자

그러면, $\text{Cost}_i = \sum\limits_{l \in [L, R], l \neq i} {|x_i - x_l|}$로 나타낼 수 있고, 고등학교 떄 학습한 것처럼 절댓값을 풀면 위 식은 $ \sum\limits_{l \in [L, x_i - 1]} {x_i - x_l} + \sum\limits_{l \in [x_i + 1, R]} {x_l - x_i}$로 나타낼 수 있다. 시그마를 처리해주면 $(n_{f} * x_i - S([L, x_i - 1])) + (S([x_i + 1, L) - n_{b} * x_i)$로 구할 수 있다. (이떄, $n_f$는 $[L, x_i - 1]$위치에 속한 나무의 개수, $n_b$는 $[x_i + 1,  R]$위치에 속한 나무의 개수이다.)

 

그러면 우리가 이 값을 구하기 위하기 위해서는

1) 쿼리에 속한 개수

2) 구간 합

을 알아야 한다.

 

위 2가지 정보 모두 세그먼트 트리로 쿼리에 대한 정보를 저장해두면 쉽게 처리할 수 있다.

Code

#include <bits/stdc++.h>
#define fastio ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);

using namespace std;
typedef long long ll;
ll seg_num[900000];
ll seg_sum[900000];


void update_num(int node_index, int node_left, int node_right, int index){
    if(index < node_left || node_right < index) return;
    seg_num[node_index]++;
    if(node_left == node_right) return;
    int mid = (node_left + node_right) / 2;
    update_num(node_index * 2, node_left, mid, index);
    update_num(node_index * 2 + 1, mid + 1, node_right, index);
    return;
}

ll query_num(int node_index, int node_left, int node_right, int query_left, int query_right){
    if(query_right < node_left || node_right < query_left) return 0;
    if(query_left <= node_left && node_right <= query_right) return seg_num[node_index];
    else{
        int mid = (node_left + node_right) / 2;
        return query_num(node_index * 2, node_left, mid, query_left, query_right) + 
               query_num(node_index * 2 + 1, mid + 1, node_right, query_left, query_right);
    }
}

ll update_sum(int node_index, int node_left, int node_right, int index, int value){
    if(index < node_left || node_right < index) return seg_sum[node_index];
    else if(node_left == node_right) return seg_sum[node_index] += value;
    int mid = (node_left + node_right) / 2;
    return seg_sum[node_index] = update_sum(node_index * 2, node_left, mid, index, value) +
                                 update_sum(node_index * 2 + 1, mid + 1, node_right, index, value);
}

ll query_sum(int node_index, int node_left, int node_right, int query_left, int query_right){
    if(query_right < node_left || node_right < query_left) return 0;
    if(query_left <= node_left && node_right <= query_right) return seg_sum[node_index];
    int mid = (node_left + node_right) / 2;
    
    return query_sum(node_index * 2, node_left, mid, query_left, query_right) +
           query_sum(node_index * 2 + 1, mid + 1, node_right, query_left, query_right);
}

int main(void){
    fastio;
    memset(seg_num, 0, sizeof(seg_num));
    memset(seg_sum, 0, sizeof(seg_sum));

    int N;
    cin >> N;

    ll total_cost = 1;
    
    for(int i = 0; i < N; i++){
        int t;
        cin >> t;
        ll front = 0;
        ll back = 0;

        if(t != 0) front = (((query_num(1, 0, 200000, 0, t - 1) % 1000000007) * t) - query_sum(1, 0, 200000, 0, t - 1)) % 1000000007 ;
        if(t != 200000) back = (query_sum(1, 0, 200000, t + 1, 200000)- ((query_num(1, 0, 200000, t + 1, 200000) % 1000000007) * t)) % 1000000007;
        ll cost = (front + back) % 1000000007;
        if(i != 0) total_cost *= cost;
        total_cost %= 1000000007;
        update_sum(1, 0, 200000, t, t);
        update_num(1, 0, 200000, t);
    }

    cout << total_cost << "\n";
    return 0;
}

구현 자체는 위에서 설명한 것처럼 처리하면 되는데, 모듈러 연산 때문에 엄청 애를 먹었다.

음수에 % 연산을 처리해주면, 기본적으로 기대한 양수 나머지가 나오지 않는다는 점을 유의하도록 하자. 

 

+ 파이썬은 잘 변환해준다고 한다.