BOJ 3653 영화 수집

문제 출처 : ICPC NWERC Regional 2011 C

 

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

 

백준에서 문제 풀다가 세그먼트 트리를 어떻게 활용할 수 있는지 보여주는 좋은 예시가 된 것 같아서 첫 개별 문제 포스팅을 정했다.

---

요점은, $n$개의 영화가 있고, 이게 순서대로 잘 꽃혀 있는데, 중간 뭔가를 빼서 보고 나면 맨 위에 놓는다. 이 행동을 $m$번 하는데, 매번 영화를 고를 때마다 위에서부터 몇 번째에 있는지를 확인하면 된다.

 

예를 들어, 1 2 3 4 5 순서로 정렬되어 있는데 여기서 4번을 빼고 나면, 4 1 2 3 5 로 바뀐다. 이때, 4번 위에는 원래 3개가 있었다. 이때, 2번을 꺼내면, 2번 위에 원래 2개가 있었고, 2 4 1 3 5 로 바뀐다. 

볼드로 처리된 값들을 출력하는 문제.

---

Naive하게 쿼리 $m$개를 처리하는 방법을 생각할 수 있다. Random Access가 가능하고, 앞에 잘 집어 넣을 수 있는 자료구조를 생각해 보면, 하나를 빼서 앞에 놓고 빈 자리를 지우는데 한 쿼리에 대략 $\mathcal{O}(n)$ 시간이 걸릴 것이기 때문에 (vector에서 중간에 하나를 빼고 그걸 다시 뒤에 넣는 식으로 구현하면, 중간에 빈 자리를 전부 당겨서 처리해야 하기 때문에 $n$ 시간이 든다), $\mathcal{O}(mn)$ 시간이 들며, $n, m = 100,000$ 까지 가능한 이 문제를 풀 수 없다.

 

즉, 한 쿼리를 적어도 $\mathcal{O}(\log n)$ 정도 시간에 처리해 내는 방법을 생각해야 하는데.. 

잘 생각해 보면, 앞 Naive Solution에서 오래 걸리는 과정은 원소들을 앞으로 당기는 과정에서 $n$개의 원소들을 처리해야 하는 경우가 발생하기 때문이다. 즉, 일단 밀지 않고 그냥 적당히 빈 자리들을 만든다고 생각하자. 즉, 앞에서 본 예시를 처리하는 데 있어서, 

1 2 3 4 5 에서 4번을 뺄 때, 4 1 2 3 0 5 로 처리하자. 그다음, 2 4 1 0 3 0 5 로 가면, $m$개 정도의 앞 공간을 쓰게 된다. 따라서, 미리 $n+m$ 개의 배열을 선언해 놓고, 1번부터 $n$번 까지의 인덱스가 지금 현재 어디에 있는지를 보관하기 위한 $n$개짜리 배열을 쓰자.

 

이렇게 하면, 

0 0 1 2 3 4 5 

2 3 4 5 6 (각 원소의 0-based index) 

위와 같이 두 배열을 들고 있는 상태에서, 4번을 빼서 앞에 놓으면

 

0 4 1 2 3 0 5

2 3 4 1 6 

이렇게 바꿔 주고, 다음 연산으로 2번을 빼면

 

2 4 1 0 3 0 5

2 0 4 1 6

이런 식으로 처리해 주면 된다.

 

이제, 사실 어차피 인덱스가 보관되어 있으면 중요한 것은 내 앞에 몇 번이 있는지는 전혀 중요하지 않고, 몇 개가 있는지만 중요하다는 사실을 알 수 있다. 따라서, 헷갈리는 실제 번호 대신, 0과 1만 사용하자. 앞 상황이 끝나고 나면 

1 1 1 0 1 0 1

2 0 4 1 6

이렇게 될 것이다. 이것만으로도 원래 정보를 전혀 잃어버리지 않았으므로 문제가 없다.

 

이제, 우리가 필요한 것은, 각 연산을 할 때마다, 0번부터 (현재 인덱스 - 1) 까지 영화가 몇 개 있는지를 빨리 세면 된다. 다시 말해서 결국은 구간 부분합을 매우 빠르게 구해야 한다 라는 말이 된다.

이런 연산을 잘 할 수 있는 자료구조로는 펜윅 트리 (Fenwick Tree) 나 세그먼트 트리 (Segment Tree) 가 있으며, 둘 다 이 문제를 풀 수 있다. $n+m$ 개의 Leaf node를 갖는 세그트리를 만들고, 이들의 합을 위로 올려가며 저장하는 식으로 트리를 구성한 다음, 업데이트를 할 때는 원하는 자리를 1에서 0으로 또는 0에서 1로 바꿔주면서 하면 된다. 

 

세그먼트 트리 구현은 여기에 매우 잘 설명되어 있고, 이걸 조금만 바꾸면 AC를 받을 수 있다.

 

#include <bits/stdc++.h>
#define usecppio ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
using namespace std;
typedef pair <int, int> pii;

struct SegTree
{
    int n;
    vector<int> segtree;

    SegTree(const vector<int> &data)
    {
        n = data.size();
        segtree.resize(4 * n);
        initialize(data, 0, n - 1, 1);
    }
    int initialize(const vector<int> &data, int l, int r, int node)
    {
        if (l == r)
            return segtree[node] = data[l];
        int mid = (l + r) / 2;
        int ls = initialize(data, l, mid, node * 2);
        int rs = initialize(data, mid + 1, r, node * 2 + 1);
        return segtree[node] = ls + rs;
    }
    void update(int ind, int diff){update(0,n-1,1,ind,diff);}
    void update(int left, int right, int node, int ind, int diff)
    {
        if (ind < left || ind > right)
            return;
        segtree[node] += diff;
        if (left != right)
        {
            int mid = (left + right) / 2;
            update(left, mid, node * 2, ind, diff);
            update(mid +1,right, node * 2 +1, ind, diff);
        }
    }
    int sum(int l, int r){return sum(l,r,1,0,n-1);}
    int sum(int l, int r, int node, int nodeleft, int noderight)
    {
        if (r < nodeleft || noderight < l)
            return 0;
        if (l <= nodeleft && noderight <= r)
            return segtree[node];
        int mid = (nodeleft + noderight) / 2;
        return sum(l,r,node*2,nodeleft,mid) + sum(l,r,node*2+1,mid+1,noderight);
    }
};



int main()
{
    usecppio
    int tc;
    cin >> tc;
    while(tc--)
    {
        int n, m;
        cin >> n >> m;
        vector <int> data(n+m,0);
        vector <int> arr(n+1, 0);
        for (int i = 1; i<=n; i++)
        {
            arr[i] = m+i-1;
            data[m+i-1] = 1;
        }
        SegTree ST(data);
        int pt = 1;
        for (int i = 0; i<m; i++)
        {
            int x;
            cin >> x;
            cout << ST.sum(0,arr[x]-1) << ' ';
            ST.update(arr[x],-1);
            arr[x] = m-pt;
            ST.update(arr[x],1);
            pt++;
        }
        cout << '\n';
    }
}