o1的题解

#include <bits/stdc++.h>
using namespace std;

// Fast IO
void fast_io() {
    ios::sync_with_stdio(false);
    cin.tie(0);
}

// Binary Indexed Tree (Fenwick Tree) for prefix sums and k-th element queries
struct BIT {
    int size;
    vector<int> tree;

    BIT(int n) : size(n), tree(n + 2, 0) {}

    // Update the BIT at position idx by delta
    void update(int idx, int delta = 1) {
        while (idx <= size) {
            tree[idx] += delta;
            idx += idx & (-idx);
        }
    }

    // Query the prefix sum up to and including idx
    int query(int idx) const {
        int res = 0;
        int x = idx;
        while (x > 0) {
            res += tree[x];
            x -= x & (-x);
        }
        return res;
    }

    // Find the k-th smallest number in the BIT
    int find_kth(int k) const {
        int idx = 0;
        int bitMask = 1;
        while (bitMask <= size) bitMask <<= 1;
        bitMask >>=1;
        while (bitMask >0) {
            int temp = idx + bitMask;
            if (temp <= size && tree[temp] < k) {
                idx = temp;
                k -= tree[temp];
            }
            bitMask >>=1;
        }
        return idx +1;
    }
};

// Function to parse the input array from the string
vector<int> parse_array(const string& s) {
    vector<int> arr;
    int start = s.find('(');
    int end = s.find(')');
    string nums = s.substr(start +1, end - start -1);
    // Split by commas
    int num = 0;
    bool negative = false;
    for(char c: nums){
        if(c == '-'){
            negative = true;
        }
        else if(c == ','){
            arr.push_back(negative ? -num : num);
            num =0;
            negative = false;
        }
        else if(isdigit(c)){
            num = num *10 + (c - '0');
        }
    }
    // Push the last number
    if(!nums.empty()){
        arr.push_back(negative ? -num : num);
    }
    return arr;
}

// Function to format and output the array
string format_output(char type, const vector<int>& arr){
    string res = "";
    res += type;
    res += "=(";
    for(int i=0;i<arr.size(); ++i){
        res += to_string(arr[i]);
        if(i != arr.size()-1) res += ",";
    }
    res += ")";
    return res;
}

int main(){
    fast_io();
    int N;
    cin >> N;
    string line;
    cin >> ws; // consume any whitespace
    getline(cin, line);
    // Determine if it's A or B
    char type = line[0];
    vector<int> arr = parse_array(line);
    if(type == 'A'){
        // Encode A to B
        vector<int> B(N, 0);
        BIT bit(N);
        // Insert A[0]
        B[0] =0;
        bit.update(arr[0]+1); // +1 because BIT is 1-indexed
        for(int i=1;i<N;++i){
            // Query number of elements less than A[i}
            if(arr[i] ==0){
                B[i] =0;
            }
            else{
                B[i] = bit.query(arr[i]);
            }
            bit.update(arr[i]+1);
        }
        // Output B
        cout << format_output('B', B);
    }
    else{
        // Decode B to A
        vector<int> A(N, 0);
        BIT bit(N);
        // Initialize BIT with all numbers available
        for(int i=0;i<N;++i){
            bit.update(i+1);
        }
        // Process from N-1 downto 0
        for(int i=N-1;i>=0; --i){
            int k = arr[i] +1;
            // Find the k-th smallest available number
            int pos = bit.find_kth(k);
            A[i] = pos -1; // since numbers are from 0 to N-1
            // Remove the number from available
            bit.update(pos, -1);
        }
        // Output A
        cout << format_output('A', A);
    }
}