백준 20921번 그렇고 그런 사이

문제 요약

정수 N, K가 주어진다. 1부터 N까지를 한 번씩 사용해서 inversion이 정확히 K개인 배열을 출력하시오.

여기서 inversion이란 배열을 $a$라고 했을 때, $i<j$인데 $a_i>a_j$인 $(i,j)$를 말한다. 이러한 $(i,j)$가 $K$개인 배열 $a$를 출력하자.

N은 최대 4242고 K는 0부터 $N(N-1)/2$이다.

풀이

일단 문제불가능한 경우는 없다고 제시했다. 그래고 한 번 왜 가능한지 생각해보자.

$a_i$가 만들 수 있는 inversion의 수는 N-i개이다. 이를 통해 0부터 N-1까지를 사용해서 그 합을 K로 만드는 것과 같다고 생각할 수 있다.

K가 0부터 N-1 사이일 때는 숫자 하나만 사용하면 된다. N부터 2N-3까지는 N-1하나랑 남은 수 중에서 하나를 선택하고, 2N-2부터 3N-6까지는 N-1, N-2 두 개와 남은 수 하나를 선택하고, ... 이런 식으로 0부터 N(N-1)/2의 수는 모두 만들 수 있다.

이런 방식을 그대로 사용해서 N-1부터 차례대로 보면서 K가 이 수보다 크다면 이 수만큼 빼주는 것을 반복해나가면 된다.

이제 N-i를 사용하는 방법이 중요한데 N-i가 필요하면 N-i+1을 배열의 제일 왼쪽부터 차례대로 채워 나가고 이 작업이 끝나면 남은 수를 오름차순으로 비어있는 곳에 배치하면 된다.

#include<bits/stdc++.h>
using namespace std;
bool vis[5005];
int ans[5005];
int main() {
    int N, K; cin >> N >> K;
    int s = 1, e = N;
    int idx = 1;
    memset(ans, -1, sizeof(ans));
    while(K) {
        if(K >= e-1) {
            K -= e-1;
            ans[idx++] = e;
            vis[e] = true;
        }
        --e;
    }
    int n = 1;
    for(int i=1;i<=N;++i) {
        if(ans[i] != -1) continue;
        while(vis[n]) ++n;
        ans[i] = n++;
    }
    for(int i=1;i<=N;++i) cout << ans[i] << " \n"[i==N];
    return 0;
}