Gorio Tech Blog search

BOJ 01280(나무 심기) 문제 풀이

11 Jul 2018 | PS Fenwick Tree BIT

목차

참조

분류 URL
문제 나무 심기
이 글에서 설명하는 코드 01280_나무 심기

개요

시간복잡도: $ O(N \log M) $

공간복잡도: $ O(M) $

  • N은 원소의 수, M은 좌표의 범위이다.

순서대로 하나씩 주어지고, $i$번째에 대한 비용이나 시간이 1부터 $(i-1)$번째까지에만 의존하는 형식이라면 BIT, IT, 삽입 정렬, 우선순위 큐를 한번쯤 떠올려 보는 것이 중요하다.

문제 풀이

우선 문제의 정의에 따라, 문제의 정답(=2번 나무부터 N번 나무까지 심는 비용의 곱)을 수식으로 정리하면 다음과 같다.

$ A := $ 문제의 정답
$ p_i := i$번 나무가 심어진 위치
$ c_i := i$번 나무를 심는 비용

$ c_i = \displaystyle\sum_{j=1}^{i-1}{\vert(p_i - p_j)\vert} $

$ A = \displaystyle\prod_{i=2}^n{c_i} $

이제 약간의 수식 변형을 시키자.

$ c_i
= \displaystyle\sum_{j=1}^{i-1}{|(p_i - p_j)|} $

$ = \displaystyle\sum_{j: p_j \le p_i, \ 1 \le j \le i-1}{(p_i - p_j)} +
\displaystyle\sum_{j: p_j \gt p_i, \ 1 \le j \le i-1}{(p_j - p_i)} $

여기서 $ n_i := p_j < p_i $인 $j$의 개수라고 하자. 그러면,

$ c_i = ( n_i \cdot p_i - \displaystyle\sum_{j: p_j \le p_i, \ 1 \le j \le i-1}{p_j} ) + ( \displaystyle\sum_{j: p_j \gt p_i, \ 1 \le j \le i-1}{p_j} - (i-1-n_i) \cdot p_i )$

$ = (2n_i+1-i)\cdot p_i + \displaystyle\sum_{j: p_j \gt p_i, \ 1 \le j \le i-1}{p_j} - \displaystyle\sum_{j: p_j \le p_i, \ 1 \le j \le i-1}{p_j} $

그러면 이제 문제는

  1. $i$번째 나무보다 왼쪽에 심어져 있던 나무의 수와
  2. $i$번째 나무보다 왼쪽에 심어져 있던 나무들의 좌표(위치)의 합

을 알아야 하는 문제로 바뀐다($p_i$와 같은 나무는 어느 쪽에 넣어도 상관없다).

  • $i$번째 나무와 위치가 같은 경우는 어차피 0이 되니 처리를 해주든 말든 상관이 없다.
  • 그리고 오른쪽에 있는 나무들은 전체 나무에서 왼쪽 부분을 빼면 얻을 수 있다(사실 그냥 오른쪽을 구해버려도 된다).

이제 펜윅 트리를 쓸 차례다. 펜윅 트리를 두 개 만든다. 하나는 위치에 따른 나무 수의 합을 세는 것이고, 하나는 위치에 따른 나무 좌표의 합을 세는 것이다. 맨 아래의 구현에서 각각 count[]와 summ[]으로 구현되어 있다.

그리고 $i$번째 나무의 위치로 $p_i$가 들어오면, 왼쪽에 있는 나무의 수는 sum($p_i$, count=true)으로, 왼쪽에 있는 나무의 좌표의 합은 sum($p_i$, count=false)로 바로 구할 수 있다.

  • 구현한 코드를 보면 이해할 수 있다. update의 코드 중복을 방지하기 위해 bool 형인 count 파라미터를 받는다.
  • 구간 합(개수 혹은 좌표의 합)을 구하는 연산은 일반 펜윅 트리와 똑같다.

그리고 업데이트를 할 때는 위치는 $p_i$로 하고 count[]는 +1, summ[]에는 $p_i$만큼 더해 준다(두 배열의 정의에 의해).

오른쪽 나무에 대해 구할 때는 전체에서 왼쪽 나무들의 값을 빼면 어렵지 않게 구할 수 있다.

구현

이번에는 펜윅 트리의 형태가 조금 다르기 때문에 따로 필자의 라이브러리를 사용하지 않았다.

주의 사항이 조금 있다.

  • 좌표가 0부터 시작하는데, 펜윅 트리는 인덱스 0을 쓰지 않는다. 따라서 $p_i$를 1씩 증가시켜서 구현해야 한다. MAX_X가 200000 + 1인 이유가 이것이다.
  • 펜윅 트리의 크기는 원소의 수가 아닌 좌표의 범위이다.
  • MOD 연산은 그때그때 해주어야 한다. 그렇다고 남용해서 마이너스를 할 때 음수가 되어버리면 연산이 제대로 된다는 보장이 없다.
  • count[]는 int[]로 선언해도 되지만, int와 long long 사이에서 실수가 많다면 그냥 편하게 모든 것을 long long으로 해도 된다. 메모리가 아주 타이트한 것이 아니라면.
#include <iostream>
#include <vector>

using namespace std;

#define ll long long
#define MOD 1000000007
#define MAX_X 200001

class FenwickTree {
public:
    vector<ll> count, summ;

    FenwickTree() {
        count.resize(MAX_X + 1);
        summ.resize(MAX_X + 1);
    }

    void update(int x, int val) {
        while (x <= MAX_X) {
            ++count[x];
            summ[x] += val;
            x += (x&-x);
        }
    }

    ll sum(int x, bool cnt) {
        ll ret = 0;
        while (x) {
            ret += cnt? count[x] : summ[x];
            x &= x-1;
        }
        return ret;
    }
};

FenwickTree BIT;
int N, p_i;

int main_01280() {
    //freopen("input.txt", "r", stdin);
    ios::sync_with_stdio(false);    cin.tie(NULL);
    cin >> N >> p_i;    ++p_i;

    BIT.update(p_i, p_i);

    ll ans = 1;
    ll p_i_sum = p_i;

    for(int i=2;i<=N;i++){
        cin >> p_i; ++p_i;

        ll n_i = BIT.sum(p_i, true);
        ll smaller = BIT.sum(p_i, false);

        ll cost = (2 * n_i + 1 - i) * p_i;
        cost -= smaller;
        cost += p_i_sum - smaller;
        cost %= MOD;

        p_i_sum += p_i;
        ans = (ans*cost) % MOD;

        BIT.update(p_i, p_i);
    }

    cout << ans;

    return 0;
}

주의: 이 코드를 그대로 복붙하여 채점 사이트에 제출하면 당연히 틀린다. 저런 헤더 파일이 채점 사이트에 있던가?

Related Posts