백준 22876 - 츠바메가에시

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

 

해당 문제는 공식 풀이에서 잘 설명하고 있으므로 이 포스트에서는 풀이 방법으로 제시된 indexed heap 풀이의 실 구현을 제공한다.

 

Indexed heap에 대해서는 여기를 참조하라.

 

전체 코드

#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
 
#include "bits/stdc++.h"
#include "ext/rope"
 
using namespace std;
using namespace __gnu_cxx;
 
using ll = long long;
using pii = pair<int, int>;
 
template <typename T>
struct Elem {
  int i;
  T w;
  bool operator < (const Elem &o) const {
    return w != o.w ? w < o.w : i < o.i;
  }
};
 
int n;
int crdx[300001];
int crdy[300001];
int score[300001];
ll arx[1000001];
ll ary[1000001];
vector<pii> gx[1000001];
vector<pii> gy[1000001];
Elem<ll> max2[2];
 
template <typename T>
struct idx_heap {
  Elem<T> *arr;
  int *pos;
  int size;
 
  idx_heap(int _n) : size(0) {
    arr = new Elem<T>[_n + 1];
    pos = new int[_n + 1];
  }
 
  ~idx_heap() {
    delete[] arr;
    delete[] pos;
  }
 
  void push(int i, T w) {
    arr[++size] = {i, w};
    pos[i] = size;
    up(size);
  }
 
  void change(int i, T w) {
    int cur = pos[i];
    auto k = arr[cur].w;
    arr[cur].w = w;
    k < w ? down(cur) : up(cur);
  }
 
  void update(int i, T w) {
    int cur = pos[i];
    auto k = arr[cur].w;
    arr[cur].w += w;
    k < arr[cur].w ? down(cur) : up(cur);
  }
 
  void up(int cur) {
    while (cur > 1) {
      if (arr[cur].w >= arr[cur / 2].w) break;
      swap(arr[cur], arr[cur / 2]);
      pos[arr[cur].i] = cur;
      cur /= 2;
    }
    pos[arr[cur].i] = cur;
  }
 
  void down(int cur) {
    while (cur * 2 <= size) {
      int mx;
      if (cur * 2 == size or arr[cur * 2].w < arr[cur * 2 + 1].w) mx = cur * 2;
      else mx = cur * 2 + 1;
      if (arr[cur].w <= arr[mx].w) break;
      swap(arr[cur], arr[mx]);
      pos[arr[cur].i] = cur;
      cur = mx;
    }
    pos[arr[cur].i] = cur;
  }
 
  T pop() {
    T ret = arr[1].i;
    arr[1] = arr[size--];
    pos[arr[1].i] = 1;
    down(1);
    return ret;
  }
 
  void del(int i) {
    int cur = pos[i];
    T k = arr[cur].w;
    arr[cur] = arr[size--];
    pos[arr[cur].i] = cur;
    k < arr[cur].w ? down(cur) : up(cur);
  }
};
 
void solve()
{
  cin >> n;
  for (int i = 1; i <= n; i++)
  {
    cin >> crdx[i] >> crdy[i] >> score[i];
    gx[crdx[i]].emplace_back(crdy[i], score[i]);
    gy[crdy[i]].emplace_back(crdx[i], score[i]);
    arx[crdx[i]] += score[i];
    ary[crdy[i]] += score[i];
  }
 
  idx_heap<ll> xpq(1e6 + 1), ypq(1e6 + 1);
 
  for (int i = 0; i <= 1e6; i++)
  {
    xpq.push(i, -arx[i]);
    ypq.push(i, -ary[i]);
  }
 
  ll ans = 0;
  ll tmp;
 
  // x loop
  for (int i = 0; i <= 1e6; i++)
  {
    tmp = arx[i];
    for (auto &item : gx[i])
    {
      auto [idx, fee] = item;
      ypq.update(idx, fee);
    }
 
    max2[0] = ypq.arr[1];
    ypq.pop();
 
    max2[1] = ypq.arr[1];
    ypq.pop();
 
    tmp += -max2[0].w - max2[1].w;
 
    ans = max(ans, tmp);
 
    ypq.push(max2[0].i, max2[0].w);
    ypq.push(max2[1].i, max2[1].w);
 
    for (auto &item : gx[i])
    {
      auto [idx, fee] = item;
      ypq.update(idx, -fee);
    }
  }
 
  // y loop
  for (int i = 0; i <= 1e6; i++)
  {
    tmp = ary[i];
    for (auto &item : gy[i])
    {
      auto [idx, fee] = item;
      xpq.update(idx, fee);
    }
 
    max2[0] = xpq.arr[1];
    xpq.pop();
 
    max2[1] = xpq.arr[1];
    xpq.pop();
 
    tmp += -max2[0].w - max2[1].w;
 
    ans = max(ans, tmp);
 
    xpq.push(max2[0].i, max2[0].w);
    xpq.push(max2[1].i, max2[1].w);
 
    for (auto &item : gy[i])
    {
      auto [idx, fee] = item;
      xpq.update(idx, -fee);
    }
  }
 
  tmp = -xpq.arr[1].w;
  xpq.pop();
  tmp -= xpq.arr[1].w;
  xpq.pop();
  tmp -= xpq.arr[1].w;
 
  ans = max(ans, tmp);
 
  tmp = -ypq.arr[1].w;
  ypq.pop();
  tmp -= ypq.arr[1].w;
  ypq.pop();
  tmp -= ypq.arr[1].w;
 
  ans = max(ans, tmp);
 
  cout << ans;
}
 
int main()
{
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
//  freopen("output.txt", "w", stdout);
#endif
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
 
  solve();
}

제출 기록