樹狀陣列之查詢中位數詳解(1057 Stack (30 分))
intuition
最近在刷pta甲級的題目, 解題過程中遇到一個之前沒有用過的知識點——樹狀陣列, 原題連結在這裡1057 Stack (30 分), 題目的大致意思是在傳統的棧的基礎上,要新增一個查詢當前棧中元素的中位數的功能, 首先想到可以用通過維護一個BST來做, 在一個二叉搜尋樹中插入和刪除元素都是資料結構的基本內容, 查詢在所有節點中比當前節點小的節點數也可以在log(n)的複雜度實現, 實現的細節在最後補充中進行了說明. 但實現一個BST資料結構做一道OJ題來說有些傻和重, 所以搜尋了一下題解, 發現可以通過樹狀陣列來得到中位數.
在介紹樹狀陣列的查詢中位數應用之前,有兩個知識點需要介紹:
- 中位數的性質:很重要 對於一個大小為N的序列中的中位數, 序列中存在(N+1) / 2 - 1個數比它小. 要找中位數也就是要統計比在序列中比當前值小的元素個數是否是(N+1)/2-1個.
- 樹狀陣列的概念和特點.
樹狀陣列
網上有不少講樹狀陣列的中文資料, 但很多都講得不夠直觀. 我好不容易搞懂了之後在這裡做個記錄.
首先這張圖反映了樹狀陣列的儲存內容特點.
樹狀陣列的元素並不都是儲存自己的值, 對於有的位置的元素, 它儲存的是一個區間的元素的和, 比如2位置儲存的是1和2的值之和, 12儲存的是9, 10, 11, 12的元素的值之和, 而8位置儲存的是1到8所有元素之和, 每個位置儲存哪個區間的值之和的規律是什麼呢?
這要講到二進位制相關的規律.
我們看下這幾個位置的索引的二進位制有什麼特點
index | 二進位制數值 | 儲存的元素之和包含哪些 | 儲存的元素之和包含的元素個數 |
---|---|---|---|
2 | 0010 | 1, 2 | 2 |
12 | 1100 | 9,10, 11,12 | 4 |
8 | 1000 | 1,2,…,8 | 8 |
回過來, 看這幾個索引的二進位制包含的零的數量k和其儲存的元素之和包含的元素個數m 存在
的關係.
也即該索引的二進位制表示中, 第一個1表示的數與其包含的求和區域有關, 比如12的第一個1
出現在100
即是4
,8的第一個1
出現在1000
也即8
. 索引和該索引位置儲存的求和範圍的關係我們搞清楚了, 接下來要回答的是, 這種關係能做什麼.
樹狀陣列的常見應用之字首和
這種關係的一個典型應用是動態字首和. 所謂動態字首和就是對於一個序列 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, 初始化為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 同樣可以得到一個值在樹中的排序, 即使這個值不在樹中. 但是要實現一整套資料結構比較厚重, 程式碼量不小也不靈活, 於是想看看能否有更簡單靈活的方式實現, 就找到了樹狀陣列.