[백준 1280] 나무심기 (C++)

1280번: 나무 심기

 

세그먼트 트리를 응용하는 문제

 

 

tree 배열은 나무를 심을 수 있는 좌표의 구간(0 ~ 200000)을 node번호로 하는 세그먼트 트리이다.

0 ~ 200000 : 1번 노드, 0 ~ 100000 : 2번 노드, 100001 ~ 200000 : 3번노드 ......

 

각 구간에 속하는 나무의 개수가 tree[node].first에 저장된다.

각 구간의 나무 위치의 합이 tree[node].second에 저장된다.

 

 

이 정보들이 왜 필요한가??

 

 

새로운 나무를 심을 때 드는 비용 구하기

나무의 좌표가 주어질 것이다.

해당 좌표보다 왼쪽 위치에 있는 나무들의 개수(first), 이들의 위치의 합(second)을 구할 수 있다.

마찬가지로 해당 좌표보다 오른쪽 위치에 있는 나무들의 개수(first), 이들의 위치의 합(second)을 구할 수 있다.

좌표와 동일한 위치에 있는 나무들은 비용 계산에서 제외된다.

 

왼쪽 나무들의 개수 * 현재 좌표 값 - 왼쪽 나무 좌표들의 합

 +

오른쪽 나무 좌표들의 합 - 오른쪽 나무들의 개수 * 현재 좌표 값

 

위 식을 통해서 현재 나무를 심는데 드는 비용을 구할 수 있다.

 

 

모두 같은 좌표가 들어올 때(set으로 체크)와

새로운 나무를 심을 때 드는 비용이 0인 경우를 예외처리 해주니 AC를 받을 수 있었다.

 

#include <iostream>
#include <set>
 
#define fastio ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
#define NUM 200001
#define MOD 1000000007
typedef long long ll;
#define pll pair<ll, ll>
 
using namespace std;
 
 
pll tree[NUM * 4];
int N;
 
void input(){
    cin >> N;
}
 
pll makeTree(int s, int e, int node, int coord){
    if(coord < s || coord > e){
        return tree[node];
    }
    if(s == e && s == coord){
        tree[node].first++;
        tree[node].second += coord;
        return tree[node];
    }
    int mid = (s + e) / 2;
    pll left = makeTree(s, mid, node*2, coord);
    pll right = makeTree(mid+1, e, node*2+1, coord);
    tree[node].first = (left.first + right.first);
    tree[node].second = (left.second + right.second);
    return tree[node];
}
 
pll leftQuery(int s, int e, int node, int coord){
    if(e < coord) return tree[node];
    if(coord <= s) return {0, 0};
    
    int mid = (s + e) / 2;
    pll left = leftQuery(s, mid, node*2, coord);
    pll right = leftQuery(mid+1, e, node*2+1, coord);
    ll first = (left.first + right.first);
    ll second = (left.second + right.second);
    return {first, second};
}
 
pll rightQuery(int s, int e, int node, int coord){
    if(s > coord) return tree[node];
    if(coord >= e) return {0, 0};
    
    int mid = (s + e) / 2;
    pll left = rightQuery(s, mid, node*2, coord);
    pll right = rightQuery(mid+1, e, node*2+1, coord);
    ll first = (left.first + right.first);
    ll second = (left.second + right.second);
    return {first, second};
}
 
void solve(){
    int x;
    ll ans = 1, la, ra;
    pll left, right;
    set<int> st;
    for(int i = 1; i <= N; i++){
        cin >> x;
        st.insert(x);
        left = leftQuery(0, NUM, 1, x);
        right = rightQuery(0, NUM, 1, x);
        la = left.first * x - left.second;
        ra = right.second - x * right.first;
        makeTree(0, NUM, 1, x);
        
        if((la + ra) == 0) continue;
        if(i != 1) ans = (((la + ra) % MOD) * ans) % MOD;
    }
    if(st.size() == 1){
        cout << 0;
    }
    else cout << ans;
}
 
int main(){
    input();
    solve();
    
    return 0;
}