CSA 079 DIV2 C K Inversions

問題概要

1 ～ N の順列で転倒数が K となるもののうち、辞書順最小のものを求めよ。

2 <= N <= 10⁵
0 <= K <= NC2

やったこと

辞書順最小ということで、とにかく今見ている部分を極力小さくする Greedy を行う。具体的には、今見ている部分の値を、「まだ使っていない数のうち k 番目に小さいものを採用したときに、可能となる最大転倒数が K 以上となる」ような最小の k を求めることを繰り返した。

このような最小の k は max(0, K - (n - i - 1)*(n - i - 2) / 2) で求められる。
このような k を求めたあとは、残った数のうちの k 番目に小さいものを BIT で管理して BIT 上二分探索によって求めた。

しかし、後でよく考えたら、BIT 上二分探索などしなくてもよかった。
「残っている数のうち一番小さい数」を採用できなくなった地点で、その段階で採用可能な最小を求めたら、あとはすべて大きい順に並べればよい話だった。

BIT を使わないコード

#include <iostream>
#include <vector>
using namespace std;

long long N, K;

int main() {
    while (cin >> N >> K) {
        vector<int> res;
        for (int i = 0; i < N; ++i) {
            long long rem = (N - i - 1) * (N - i - 2) / 2;
            if (rem >= K) res.push_back(i + 1);
            else {
                int kth = K - rem;
                res.push_back(i + 1 + kth);
                for (int j = N - 1; j >= i; --j) {
                    if (j == i + kth) continue;
                    res.push_back(j + 1);
                }
                break;
            }
        }
        for (int i = 0; i < (int)res.size(); ++i) {
            cout << res[i];
            if (i != (int)res.size() - 1) cout << " ";
        }
        cout << endl;
    }
}

BIT を使ったコード

#include <iostream>
#include <vector>
using namespace std;

template <class Abel> struct BIT {
    vector<Abel> dat;
    Abel UNITY_SUM = 0;                        // to be set

    /* [1, n] */
    BIT(int n) { init(n); }
    void init(int n) {
        dat.resize(n + 1);
        for (int i = 0; i < (int)dat.size(); ++i) dat[i] = UNITY_SUM;
    }

    /* a is 1-indexed */
    inline void add(int a, Abel x) {
        for (int i = a; i < (int)dat.size(); i += i & -i)
            dat[i] = dat[i] + x;
    }

    /* [1, a], a is 1-indexed */
    inline Abel sum(int a) {
        Abel res = UNITY_SUM;
        for (int i = a; i > 0; i -= i & -i)
            res = res + dat[i];
        return res;
    }

    /* [a, b), a and b are 1-indexed */
    inline Abel sum(int a, int b) {
        return sum(b - 1) - sum(a - 1);
    }

    /* k-th number (k is 0-indexed) */
    int get(long long k) {
        ++k;
        int res = 0;
        int N = 1; while (N < (int)dat.size()) N *= 2;
        for (int i = N / 2; i > 0; i /= 2) {
            if (res + i < (int)dat.size() && dat[res + i] < k) {
                k = k - dat[res + i];
                res = res + i;
            }
        }
        return res + 1;
    }
};

long long N, K;

int main() {
    while (cin >> N >> K) {
        vector<int> res;
        BIT<int> bit(N);
        for (int i = 1; i <= N; ++i) bit.add(i, 1);

        for (int i = 0; i < N; ++i) {
            long long m = N - i - 1;
            long long kth = max(0LL, K - m*(m - 1) / 2);
            int val = bit.get(kth);
            res.push_back(val);
            bit.add(val, -1);
            K -= kth;
        }

        for (int i = 0; i < res.size(); ++i) {
            cout << res[i];
            if (i != (int)res.size() - 1) cout << " ";
        }
        cout << endl;
    }
}