244. 谜一样的牛（树状数组、权值线段树）

传送门

题意

有 $n$ 头奶牛，已知它们的身高为 $1$ ~ $n$ 且各不相同，但不知道每头奶牛的具体身高。

现在这 $n$ 头奶牛站成一列，已知第 $i$ 头牛前面有 $A_i$ 头牛比它低，求每头奶牛的身高。

---

输入格式

第 $1$ 行：输入整数 $n$ 。

第 $2$ ~ $n$ 行：每行输入一个整数 $A_i$ ，第 $i$ 行表示第 $i$ 头牛前面有 $A_i$ 头牛比它低。
（注意：因为第 $1$ 头牛前面没有牛，所以并没有将它列出）

---

输出格式

输出包含 $n$ 行，每行输出一个整数表示牛的身高。

第 $i$ 行输出第 $i$ 头牛的身高。

---

数据范围

$1<=n<=10^5$

---

输入样例：

5
1
2
1
0

输出样例：

2
4
5
3
1

---

思路

• 从后向前遍历，能够确定当前牛的排名，$a[i]+1$。
• 由于我们已经确定了 $[i+1,n]$ 头牛的位置。所以在确定第 $i$ 头牛的位置时，考虑的是区间 $[1,n]$ 中除去后面这些牛的第 $a[i]+1$ 小的位置。

---

代码

• 树状数组

  • 二分确定位置

#include <iostream>
#include <cstring>
#include <algorithm>

using namespace std;

const int N = 100010;

int tr[N], a[N], b[N];
int n;

void add(int x, int c) {
    for (; x <= N; x += x & -x) tr[x] += c;
}

int sum(int x) {
    int res = 0;
    for (; x; x -= x & -x) res += tr[x];
    return res;
}

int main()
{
    cin >> n;
    add(1, 1);
    for (int i = 2; i <= n; i ++ ) cin >> a[i], add(i, 1);
    
    for (int i = n; i >= 1; i--) {
        int l = 1, r = n;
        while(l < r) {
            int mid = l + r >> 1;
            if(sum(mid) >= a[i] + 1) r = mid;
            else l = mid + 1;
        }
        b[i] = l;
        add(l, -1);
    }
    for (int i = 1; i <= n; i ++ ) cout << b[i] << endl;
    return 0;
}

---

• 权值线段树

  • query：即等同于二分操作

#include <iostream>
#include <cstring>
#include <algorithm>

#define ls u<<1
#define rs u<<1|1


using namespace std;

const int N = 100010;

int n;
int a[N], b[N];
int tr[N << 2], L[N << 2], R[N << 2];

void pushup(int u) {
    tr[u] = tr[ls] + tr[rs];
}

void build(int u, int l, int r) {
    L[u] = l, R[u] = r;
    if(l == r) { tr[u] = 1; return ;}
    int mid = l + r >> 1;
    build(ls, l, mid), build(rs, mid + 1, r);
    pushup(u);
}

void query(int u, int k, int &res) {
    int ll = L[u], rr = R[u];
    if(ll == rr) {tr[u] = 0; res = /*n - ll + 1*/ ll; return ; }
    if(tr[ls] >= k) query(ls, k, res);
    else query(rs, k - tr[ls], res);
    pushup(u);
}

int main()
{
    cin >> n;
    for (int i = 2; i <= n; i++) cin >> a[i];
    
    build(1, 1, n);
    
    // int cnt = 0; // 当前选出了多少人
    for (int i = n; i >= 1; i--) {
        query(1, /*n - cnt -  a[i]*/ a[i] + 1, b[i]); 
        // cnt++;
    }
    for (int i = 1; i <= n; i++) cout << b[i] << endl;
    return 0;
}