나무 심기

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

무난하게 구간 쿼리(segment tree or Fenwick)을 통해 풀 수 있는 문제이다.

나는 개인적으로 Fenwick 쪽이 구현이 간단해서 이를 사용하였다.

i번 나무를 심는 단계에 이르렀다고 생각해보자.

결국 절대값의 성질로 부터 i번 나무보다 앞에 위치하는 나무들에 대해서 그 거리의 합은 아래와 같다.

• i번 나무보다 앞에 위치하는 나무들과의 거리합 = SUM(i번 나무 위치 - j번 나무 위치)
  (이때, j번 나무는 i번 나무보다 앞에 존재)

위의 식을 다시 쓰면 아래와 같다.

• i번 나무보다 앞에 위치하는 나무들과의 거리합 = i번 나무 위치 * i번 보다 앞에 있는 나무의 수 - i번 보다 앞에 있는 나무의 위치 합
  (이때, j번 나무는 i번 나무보다 앞에 존재)

따라서, 핵심이 되는 항은 i번 보다 앞에 있는 나무의 수, i번 보다 앞에 있는 나무의 위치 합, 2개의 항이다.

전자의 경우 위치 배열에 나무가 심어질 때마다 1을 증가시키는 구간 트리로, 후자의 경우 위치 배열에 나무가 심어질 때마다 나무 위치의 값만큼을 증가시키는 구간트리로 쉽게 구해줄 수 있다.

i번 나무보다 뒤에 위치하는 나무들 같은 경우에도 유사하게 고려가능하다.

한가지 난 몇번 실수 했던 것이 있는데, A * B mod N == (A mod N) * (B mod N) mod N임을 잊지 말도록 하자.

#include <iostream>

using namespace std;

static const int64_t kMaxX = 200000 + 1;
static const int64_t kMod = 1000000007;

class FenwickTree
{
private:
  int64_t node_[kMaxX + 1];

public:
  FenwickTree(void)
  {
    for (int64_t it = 0; it < kMaxX + 1; ++it)
    {
      node_[it] = 0;
    }
  }

  void Update(const int64_t idx, const int64_t value)
  {
    int64_t index = idx;

    while (true)
    {
      node_[index] += value;
      index += index & (-index);

      if (index > kMaxX)
      {
        break;
      }
    }
  }

  int64_t Psum(const int64_t idx)
  {
    int64_t index = idx;

    int64_t ret = 0;

    while (true)
    {
      ret += node_[index];
      index -= index & (-index);

      if (index <= 0)
      {
        break;
      }
    }

    return ret;
  }
};

FenwickTree val_tree;
FenwickTree cnt_tree;

int main(void)
{
  // For faster IO
  ios_base::sync_with_stdio(false);
  cout.tie(nullptr);
  cin.tie(nullptr);

  // Solve
  int64_t N;
  cin >> N;

  int64_t first;
  cin >> first;

  val_tree.Update(first + 1, first + 1);
  cnt_tree.Update(first + 1, 1);

  int64_t ret = 1;
  for (int64_t it = 1; it < N; ++it)
  {
    int64_t number;
    cin >> number;

    int64_t b_val = val_tree.Psum(kMaxX) - val_tree.Psum(number);
    int64_t s_val = val_tree.Psum(number);

    int64_t b_cnt = cnt_tree.Psum(kMaxX) - cnt_tree.Psum(number);
    int64_t s_cnt = cnt_tree.Psum(number);

    int64_t dist = s_cnt * (number + 1) - s_val + b_val - b_cnt * (number + 1);
    ret = (ret * (dist % kMod)) % kMod;

    val_tree.Update(number + 1, number + 1);
    cnt_tree.Update(number + 1, 1);
  }

  cout << ret << "\n";

  return 0;
}