백준 30515 - 방형구 탐색 (Hard) (C++)

문제

• 문제 링크

  
   BOJ 30515 - 방형구 탐색 (Hard) 
  

• 문제 요약

$1*N$ 크기의 격자에서, 각 칸에 피어있는 꽃의 종류 $A_1, A_2, ... , A_N$가 주어진다.
아래 두 쿼리를 수행하는 프로그램을 작성하자.



$1$ $l$ $r$ $k$ : $[l, r]$에서 꽃의 종류가 $k$인 꽃의 개수를 출력한다.

$2$ $l$ $r$ : $[l, r]$에 속한 모든 꽃들을 제거한다.

• 제한

TL : $1.5$ sec, ML : $1024$ MB




$1 ≤ N, Q ≤ 200,000$

$1 ≤ A_i$$, k ≤ 10^9$

$1$번 쿼리는 하나 이상 주어진다.

알고리즘 분류

• 자료 구조 (data_structures)
• 트리를 사용한 집합과 맵 (tree_set)

풀이

$EASY$버전을 나이브하게 짜고 이걸 보며 어떻게 최적화할 지 고민하고 있었다.

그런데 문득, 얼마 전 이 문제 에서 써먹었던 $PBDS$가 번뜩였다.

이를 쓸 수만 있다면 기본 $set$에서 특정 요소보다 작은 요소의 개수를 $O(logN)$에 구할 수 있기 때문이다.

$A_1, A_2, ... , A_N$에서 등장하는 수의 종류만큼 $PBDS$를 만들자.

오름차순으로 트리를 구성해야 하므로 $Policy$는 $less$로 한다.

$A_i = k$인 모든 $i$들은 $PBDS_k$에 추가된다.

$order$_$of$_$key(x)$를 이용해 $PBDS_k$에서 $[l, r]$에 속한 $i$의 개수를 센다고 해보자.

일반적인 상황에선 $order$_$of$_$key(r)$ $-$ $order$_$of$_$key(l)$ $+$ $1$ 로 간단히 정리할 수 있다.

$iterator$를 다룸에 있어서 런타임 에러를 맞지 않도록 적절히 컨트롤해 주어야 할 것이다.

---

한편, 삭제할 때마다 $[l, r]$을 순회해야 할텐데 이를 매번 하나하나 순회 한다면 $O(QN)$의 복잡도를 갖게 된다.

$2$번 쿼리에서 한 번 삭제된 요소는 다음부터 전혀 고려되지 않음을 알 수 있을 것이다.

이에 따라 현재 유효한 꽃들만 갖고 있는다면 불필요한 탐색을 막을 수 있다.

최악의 경우를 산정해도, 들고 있는 집합의 크기가 점차 감소하고 필요한 만큼만 탐색한다.

즉, $amortized$ $O((N+Q)logN)$의 시간 복잡도를 갖는다.


※ 평방 분할이나, 세그먼트 트리를 이용한 온라인 / 오프라인 풀이도 있다고 한다.

전체 코드

#include<bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
 
using namespace std;
using namespace __gnu_pbds;
 
typedef tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update> PBDS;
 
int main()
{
    ios_base::sync_with_stdio(0); cin.tie(0);
    int n; cin >> n;
 
    unordered_map <int, PBDS> S;
    map <int, int> M;
    for (int i(1); i <= n; i++)
    {
        int k; cin >> k;
        S[k].insert(i);
        M[i] = k;
    }
 
    int q; cin >> q;
    while (q--)
    {
        int op, l, r; cin >> op >> l >> r;
        if (op & 1)
        {
            int k; cin >> k;
            if (!S.count(k) || !S[k].size()) { cout << "0\n"; continue; }
 
            auto it(S[k].lower_bound(l));
            auto _it(S[k].upper_bound(r));
            if (it == S[k].end() || _it == S[k].begin())
            {
                cout << "0\n";
                continue;
            }
 
            int _l(S[k].order_of_key(*it));
            int _r(S[k].order_of_key(*--_it));
 
            cout << _r - _l + 1 << '\n';
        }
        else
        {
            stack <int> st;
            for (auto it(M.lower_bound(l)); it != M.end() && (*it).first <= r; it++)
            {
                auto [idx, _k](*it);
                S[_k].erase(idx);
                st.push(idx);
            }
            while (st.size()) M.erase(st.top()), st.pop();
        }
    }
}

comment

끝