[ICPC] 백준 5012 - 불만 정렬

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

 

문제를 해결하기 위해 펜윅 트리에 다음 정보를 관리할 것이다.

 

현재 순회 중인 index에 해당하는 값 val 이상의 값이 출현하는 횟수

 

그러면 val을 pivot으로 할 때 구간은 다음과 같이 표현될 수 있다.

 

[val보다 작은 index에서 val 이상의 값이 나타나는 횟수), val, (val 보다 큰 index에서 val 이상의 값이 나타나는 횟수]

 

이렇게 되면 왼쪽 구간에서 큰 수를 선택하는 경우와 오른쪽 구간에서 작은 수를 선택하는 경우를 곱셈법칙으로 계산하여 왼쪽 구간의 값 * (오른쪽 원소의 갯수 - 오른쪽 구간의 값)를 각 pivot마다 더해주어 답을 구할 수 있다.

 

하지만 이렇게 하면 왼쪽 구간에서 val보다 큰 값이 있으리라는 보장을 할 수 없다.

 

큰 수부터 내림차순으로 펜윅 트리에 반영해주면 왼쪽 구간에서는 val보다 큰 원소만 존재함을 보장할 수 있다.

 

원소의 중복 반영을 방지하기 위해 같은 값의 원소에 대해 index의 내림차순으로 계산해줘야 올바른 답을 계산할 수 있다.

 

내림차순으로 계산을 하는 구현은 밑의 코드를 참조한다.

 

전체 코드

#pragma GCC target("avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

int n;
int maxn;
int fenwick[100001];
vector<int> idx[100001];

int sum(int pos)
{
    int ret = 0;

    while (pos)
    {
        ret += fenwick[pos];
        pos &= (pos - 1);
    }

    return ret;
}

void add(int pos, int val)
{
    while (pos <= n)
    {
        fenwick[pos] += val;
        pos += (pos & -pos);
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    cin >> n;
    for (int i = 1; i <= n; ++i)
    {
        int x;
        cin >> x;

        idx[x].push_back(i);
        maxn = max(maxn, x);
    }

    ll res = 0;

    for (int i = maxn; i >= 1; --i)
    {
        if (idx[i].empty())
            continue;
        
        for (int j = idx[i].size() - 1; j >= 0; --j)
        {
            add(idx[i][j], 1);
            int inv_right = (n - idx[i][j]) - (sum(n) - sum(idx[i][j]));
            int inv_left = sum(idx[i][j] - 1);
            res += (ll)inv_left * inv_right;
        }
    }

    cout << res;
}