백준 1572 - 중앙값

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

 

전형적인 세그먼트 트리로 k번째 수를 찾는 쿼리를 이용하는 문제이다.

 

하지만 그 구현은 너무나도 귀찮다. C++을 사용하고 있다면 pbds를 이용해 간단하게 풀 수 있다.

 

find_by_order와 order_by_key를 잘 이용하자.

 

전체 코드

#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"
 
#include <ext/pb_ds/assoc_container.hpp>
 
using namespace std;
using namespace __gnu_cxx;
using namespace __gnu_pbds;
 
using ll = long long;
using pii = pair<int, int>;
 
int n, k;
ll ans;
int ar[250000];
 
using ordered_set = tree<
    int,
    null_type,
    less_equal<>,
    rb_tree_tag,
    tree_order_statistics_node_update>;
 
void solve()
{
  cin >> n >> k;
 
  ordered_set st;
 
  for (int i = 0; i < n; i++)
  {
    int x;
    cin >> x;
    ar[i] = x;
    st.insert(x);
 
    if (st.size() > k)
    {
      int target = ar[i - k];
      auto idx = st.order_of_key(target);
      st.erase(st.find_by_order(idx));
    }
 
    if (st.size() == k)
    {
      int mid = (k + 1) / 2 - 1;
      ans += *st.find_by_order(mid);
    }
  }
 
  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();
}

제출 기록