intuition

最近在刷pta甲級的題目, 解題過程中遇到一個之前沒有用過的知識點——樹狀陣列, 原題連結在這裡1057 Stack （30 分), 題目的大致意思是在傳統的棧的基礎上,要新增一個查詢當前棧中元素的中位數的功能, 首先想到可以用通過維護一個BST來做, 在一個二叉搜尋樹中插入和刪除元素都是資料結構的基本內容, 查詢在所有節點中比當前節點小的節點數也可以在log(n)的複雜度實現, 實現的細節在最後補充中進行了說明. 但實現一個BST資料結構做一道OJ題來說有些傻和重, 所以搜尋了一下題解, 發現可以通過樹狀陣列來得到中位數.

在介紹樹狀陣列的查詢中位數應用之前,有兩個知識點需要介紹:

1. 中位數的性質:很重要 對於一個大小為N的序列中的中位數, 序列中存在(N+1) / 2 - 1個數比它小. 要找中位數也就是要統計比在序列中比當前值小的元素個數是否是(N+1)/2-1個.
2. 樹狀陣列的概念和特點.

樹狀陣列

網上有不少講樹狀陣列的中文資料, 但很多都講得不夠直觀. 我好不容易搞懂了之後在這裡做個記錄.
首先這張圖反映了樹狀陣列的儲存內容特點.

樹狀陣列的元素並不都是儲存自己的值, 對於有的位置的元素, 它儲存的是一個區間的元素的和, 比如2位置儲存的是1和2的值之和, 12儲存的是9, 10, 11, 12的元素的值之和, 而8位置儲存的是1到8所有元素之和, 每個位置儲存哪個區間的值之和的規律是什麼呢?
這要講到二進位制相關的規律.
我們看下這幾個位置的索引的二進位制有什麼特點

回過來, 看這幾個索引的二進位制包含的零的數量k和其儲存的元素之和包含的元素個數m 存在 $m$

樹狀陣列的常見應用之字首和

這種關係的一個典型應用是動態字首和. 所謂動態字首和就是對於一個序列 a1, a2, ... , an有查詢和修改操作, 查詢的內容是對於給定的am,要查詢a1+a2+...+am的值,;修改是這個序列會動態變化, 增加元素, 刪除元素, 修改元素. 對於暴力演算法, 一次修改的複雜度是O(1), 一次查詢的複雜度是O(n), 我們想要通過某種方法降低查詢的複雜度, 這就是樹狀陣列.

樹狀陣列如何查詢

由上面的那張圖可以看出, 例如要查詢13位置的字首和, 就只需要訪問位置13, 12, 8, 對他們的值累加, 就得到了13位置的元素的字首和.
我們來看二進位制視角下, 這一次查詢訪問的位置有什麼特點:

13: 1101
12: 1100
8:   1000
1101 = 1000 + 0100 + 0001

可以看出來, 訪問的位置數量就是該索引中1出現的數量, 訪問的位置和1的位置有關.

訪問13
13: 第一個1為1
13 - 1 -> 12
訪問12
12: 第一個1為100
12-4(100) -> 8
訪問8
8: 第一個1為1000
8-8 -> 0
結束

樹狀陣列如何修改

修改某個索引的值 am時, 要修改所有包含了am的元素, 對於上圖, 修改13時, 要修改14, 16, 修改9時, 要修改10, 12, 16
我們在二進位制視角下看看訪問的這些位置有什麼特點

修改a9
9: 1001 第一個1為1
9+1->10
修改10
10: 1010 第一個1為10
10+2(10)->12
修改12
12: 1100 第一個1為100
12+4(100)->16
修改16, 到達陣列的最大尺寸,結束

講完了動態字首和, 應該大家就很容易想到, 找中位數的核心是找序列中比它小的數的個數, 也可以用樹狀陣列來做. 此時當am這個值在序列中時, 我們就令am=1, 否則am=0, 統計a_q字首和就是統計比a_q小的數有多少個.

在講找中位數之前, 還要講的一個問題是如何找到一個數的二進位制表示下的第一個1的位置
回顧一個問題, 一個正數的補碼如何表示, 例如,對於6(0110),它的補碼是其反碼加1, 也就是(1001+1= 1010), 求一個二進位制數的補碼, 就是找到這個數中的第一個1, 其右邊保持不變, 左邊取反. 於是我們可以用一個很巧妙的方法, 取到第一個1的位置

int lowbit(int i) {
    return i & (-1);
}

通過該二進位制數和其補碼做與運算,保留了第一個1以及其右邊的部分, 得到的就是這個1所在的位置, 以6為例,0110 & 1010 = 0010, 即是2, 我們可以看下,6位置儲存的是6-2+1, 6-2+2,兩個值, 對於樹狀陣列的任意位置, 其儲存的正是 $a_{m-2^k+1}, a_{m-2^k +2}, ... ,a_{m-2^k+2^k}$元素之和.

利用樹狀陣列查詢中位數

按照上面所說的, 開一個足夠大的陣列a, 初始化為0, 若m值在陣列中, 我們則需要更新a[m]位置的值,令其+1, 同時修改受影響的所有位置的值, 要m值離開了陣列, 同樣要更新a[m]位置的值, 當查詢m在序列中的升序排名時, 返回的是a[m]位置的字首和.
好了, 祭上柳神的程式碼

#include <iostream>
#include <stack>
#define lowbit(i) ((i) & (-i))
const int maxn = 100010;
using namespace std;
int c[maxn];
stack<int> s;
void update(int x, int v) {
    for(int i = x; i < maxn; i += lowbit(i))
        c[i] += v;
}
int getsum(int x) {
    int sum = 0;
    for(int i = x; i >= 1; i -= lowbit(i))
        sum += c[i];
    return sum;
}
void PeekMedian() {
    int left = 1, right = maxn, mid, k = (s.size() + 1) / 2;
    while(left < right) {
        mid = (left + right) / 2;
        if(getsum(mid) >= k)
            right = mid;
        else
            left = mid + 1;
    }
    printf("%d\n", left);
}
int main() {
    int n, temp;
    scanf("%d", &n);
    char str[15];
    for(int i = 0; i < n; i++) {
        scanf("%s", str);
        if(str[1] == 'u') {
            scanf("%d", &temp);
            s.push(temp);
            update(temp, 1);
        } else if(str[1] == 'o') {
            if(!s.empty()) {
                update(s.top(), -1);
                printf("%d\n", s.top());
                s.pop();
            } else {
                printf("Invalid\n");
            }
        } else {
            if(!s.empty())
                PeekMedian();
            else
                printf("Invalid\n");
        }
    }
    return 0;
}

出處: https://github.com/liuchuo/PAT/blob/4fb16451ff/AdvancedLevel_C%2B%2B/1057. Stack (30) .cpp

補充 在BST中查詢節點的排序

public int rank(Key key) {
        if (key == null) throw new IllegalArgumentException("the argument to rank() is null.");
        return rank(root, key); // 從樹根開始搜尋
    }

    private int rank(Node x, Key key) {
        if (x == null) return 0;
        int cmp = key.compareTo(x.key);
        if (cmp > 0) return size(x.lchild) + 1 + rank(x.rchild, key); // 如果當前節點的值小於key, 則返回當前節點左子樹的大小以及該節點的計數1以及key在其右子樹中的排名.
        else if (cmp < 0) return rank(x.lchild, key); // 如果當前節點的值大於key, 則返回key在當前節點的左子樹中的排名
        else return size(x.lchild); // 如果當前節點的值與key相等, 返回其左子樹的大小.
    }

通過BST 同樣可以得到一個值在樹中的排序, 即使這個值不在樹中. 但是要實現一整套資料結構比較厚重, 程式碼量不小也不靈活, 於是想看看能否有更簡單靈活的方式實現, 就找到了樹狀陣列.