區間查詢(樹狀陣列和線段樹)
本文是LeetCode 307。 Range Sum Query - Mutable的題解,主要是對樹狀陣列(Binary Indexed Tree)和線段樹(Segment Tree)的學習。
本文首發於我的部落格,傳送門
1。 問題描述
初始時給定一個大小為n的整數陣列nums,設計兩個方法:
sumRange(i,j)
: 求區間[i, j]之間的元素和(0 <= i <= j < n),包含邊界;
update(i, val)
: 將nums[i]的值修改為val。
而且假設了上述兩個方法的呼叫是均勻分佈的。
2。 題解
2。1 初步分析
如果陣列nums是不可變的,那就是題目303。 Range Sum Query - Immutable,是一個很簡單的題,我們只需要維護一個累加陣列sum, sum[i]表示前i個元素和, 這樣以後就可以以常數複雜度獲得某個區間的和:
sumRange(i,j) = sum[j] - sum[i-1]
。如果此題也採用此思路的話,每次呼叫
update(i, val)
都需要更新sum, 即
update(i, val)
的複雜度是O(n);若採用每次呼叫
sumRange(i,j)
才進行求和的策略,雖然
update(i, val)
複雜度降成常數,但是
sumRange(i,j)
複雜度變成了O(n)。由於題目說了
sumRange
和
update
的呼叫是均勻分佈的,因此上述兩種策略在分別m次呼叫的情況下時間複雜度都是O(mn),不夠高效。
2。2 解法一、折中策略
2。2。1 分析
2。1節的分析中的兩種策略要麼讓
sumRange
為O(1)而
update
為O(n),要麼讓
sumRange
為O(n)而
update
為O(1),即兩種策略都是不公平的。所以我們很容易想到將複雜度均攤到
sumRange
和
update
上,即採取兩種策略的折中策略:將原陣列分為若干塊,為了均攤複雜度,我們讓每塊長度
block = sqrt(n)
。然後我們開闢一個大小跟塊數(=sqrt(n) 向上取整)相同的二維陣列,採用303題的策略提前計算好每個塊內部的累加和陣列。對於
sumRange
求區間和操作,我們將i和j之間塊的和累加起來(需要考慮兩端不足整塊的部分)就可以了,由於塊數是
sqrt(n)
(向上取整),所以複雜度是O(sqrt(n))。而在呼叫
update
更新的時候,我們先確定要更新的元素在哪個塊裡,然後只更新該塊對應的累加和陣列,因為塊長度最大為sqrt(n),所以複雜度還是O(sqrt(n))。因此
sumRange
和
update
的複雜度均為
O(sqrt(n))
,即在分別m次呼叫的情況下時間複雜度為O(m*sqrt(n))。
2。2。2 程式碼
class
NumArray
{
int
len
,
block
;
// nums長度, 塊長度
vector
<
vector
<
int
>>
block_sum
;
public
:
NumArray
(
vector
<
int
>&
nums
)
{
len
=
nums
。
size
();
block
=
int
(
sqrt
(
len
));
vector
<
int
>
sum
;
int
accum
=
0
;
for
(
int
i
=
0
;
i
<
len
;
i
++
){
accum
+=
nums
[
i
];
sum
。
push_back
(
accum
);
if
((
i
+
1
)
%
block
==
0
||
i
==
len
-
1
){
block_sum
。
push_back
(
sum
);
accum
=
0
;
sum
。
clear
();
}
}
}
void
update
(
int
i
,
int
val
)
{
int
block_i
=
i
/
block
,
block_j
=
i
%
block
;
// 塊號, 塊內索引
int
diff
=
val
-
block_sum
[
block_i
][
block_j
];
if
(
block_j
!=
0
)
diff
+=
block_sum
[
block_i
][
block_j
-
1
];
for
(
int
j
=
block_j
;
j
<
block_sum
[
block_i
]。
size
();
j
++
)
block_sum
[
block_i
][
j
]
+=
diff
;
}
int
sumRange
(
int
i
,
int
j
)
{
int
block_i
=
i
/
block
,
block_j
=
j
/
block
;
// 塊號, 塊號
int
res
=
0
;
for
(
int
k
=
block_i
;
k
<
block_j
;
k
++
)
res
+=
block_sum
[
k
]。
back
();
// 減去左端多加的部分
if
(
i
%
block
!=
0
)
res
-=
block_sum
[
block_i
][
i
%
block
-
1
];
// 加上右端不足整塊的部分
res
+=
block_sum
[
block_j
][
j
%
block
];
return
res
;
}
};
2。3 解法二、樹狀陣列
2。3。1 樹狀陣列(Binary Indexed Tree)
樹狀陣列(Binary Indexed Tree, 又Fenwick Tree)其實並不是一棵樹,只是對陣列各元素進行邏輯上的劃分。根據維基百科,樹狀陣列是一種用於高效計算數列字首和的資料結構,它可以以O(logn)的時間得到任意字首和(兩個字首和相減即可得到區間和),並同時支援以O(logn)的時間對陣列某個值進行修改,空間複雜度為O(n)。由此可見,我們可以用樹狀陣列解此題,使
sumRange
和
update
的複雜度均為O(logn)。因此我們有必要來了解一下樹狀陣列。
從英文名字Binary Indexed Tree可猜到樹狀陣列是與二進位制有關的,事實上確實是這樣,樹狀陣列的核心思想就是將一個字首和劃分成多個子序列的和,而劃分的方法與2的冪(或者說二進位制)密切相關。
作為對比,本文講的三個解法都是對陣列先劃分成子序列,只是劃分的方式不同。2。2節的思路就是無腦將陣列均分成大小為sqrt(n)的塊;樹狀陣列則採用精心設計的劃分方式;而2。4節講的線段樹則是不斷進行二分。
舉個例子說明這種劃分方式,例如想求前13個數(下標從1開始,下同)的字首和
sum(13)
,我們可以先看不大於13的最大的2的冪是多少,是2^3=8, 所以先分成第1到第8個數的和;還剩下5個數,又看不大於5的最大的2的冪是多少,是2^2=4,所以再分成第9到第12個數的和;最後還剩下一個數,即第13個數。所以我們有:
sum(13) = Range(1, 8) + Range(9, 12) + Range(13, 13)
其中
Range(i,j)
就是閉區間[i,j]的區間和。如果按照這種方式劃分的話,很顯然
Range(i,j)
中的引數i是可以根據j確定的,例如若j=12那麼i一定等於9,我們把這個對映關係記為f(j)。如果我們將上面的數字都轉成二進位制然後稍加分析就可得到
f(j) = j - lowBit(j) + 1;
其中
lowBit(j)
表示將j轉為二進位制後, 最後一個1的位置所代表的數值。例如
j=12 ——-二進位制——> 1100 ——-lowBit——> 0100 ——-f(j)——> 1001
學過補碼錶示的童鞋可以想到
lowBit(j)
可以透過將j與上-j得到,C++程式碼即
int
lowBit
(
int
j
){
return
j
&
(
-
j
);
}
那麼,我們用一個數組bit(Binary Indexed Tree)來表示經過我們精心設計的劃分後的子序列的區間和,其中
bit[i]
表示區間
[f(i), i]
的和,如下圖所示。
那麼bit陣列計算好後(先不管如何計算),那麼我們該如何計算字首和以及修改某個元素呢?
1. 計算字首和
由前面計算字首和
sum(13)
的例子
sum(13) = Range(1, 8) + Range(9, 12) + Range(13, 13)
= bit[8] + bit[12] + bit[13]
= bit[13] + bit[12] + bit[8]
可總結出求字首和的演算法:
int
sum
(
int
i
){
// 前idx個數的和, logn
// 這裡的i從1開始編號
int
res
=
0
;
while
(
i
>
0
){
res
+=
bit
[
i
];
i
-=
lowBit
(
i
);
}
return
res
;
}
可見可以在log(n)的複雜度求得某個字首和。
2. 更新元素
由前面的圖可知,更新某個元素也是取樣類似的思路,只是說是個逆向過程。例如如果想更新nums[1],對應下標加1為2,由上圖可知需要自底向上更新bit[2]、bit[4]、bit[8];又例如更新nums[4],對應下標加1為5,應該自底向上更新bit[5]、bit[6]、bit[8]。所以們可總結更新演算法如下
void
update
(
int
i
,
int
val
)
{
//logn
int
diff
=
val
-
nums
[
i
];
nums
[
i
]
=
val
;
i
+=
1
;
// bit中下標從1開始
while
(
i
<=
nums
。
size
()){
bit
[
i
]
+=
diff
;
i
+=
lowBit
(
i
);
}
}
可見
求字首和函式sum和更新函式update的唯一不同就是前者是不斷減去lowBit(i)而後者是不斷加上lowBit(i)
。
3. 建樹
前面的分析都是建立在我們已經計算好了陣列bit的基礎上,那麼如何計算bit呢?最簡單的思路就是先用0初始化bit,然後呼叫n次
update
就可以了,此時建樹時間複雜度為O(nlogn)。此外,存在一個O(n)時間計算bit的演算法
將
bit[i+1]
初始化成
nums[i]
(因為bit下標從1開始);
對1到n的每一個i,令
j = i + lowBit(i)
,更新
bit[j] = bit[j] + bit[i]
。
綜上,我們可以用樹狀陣列以對數時間複雜度求字首和以及更新元素,而兩個字首和相減即可得到區間和,因此
sumRange
和
update
的複雜度均為
O(logn)
,初始建樹的時間複雜度最低只需O(n);我們需要一個bit陣列,因此空間複雜度為O(n)。
2。3。2 程式碼
綜合前面的分析,我們用如下程式碼(陣列data是nums的複本):
class
NumArray
{
private
:
vector
<
int
>
bit
;
// bit中下標從1開始, bit[0] = 0
vector
<
int
>
data
;
// copy of nums
int
lowBit
(
int
x
){
return
x
&
(
-
x
);
}
int
sum
(
int
idx
){
// logn
// 前idx個數的和
int
res
=
0
;
while
(
idx
>
0
){
res
+=
bit
[
idx
];
idx
-=
lowBit
(
idx
);
}
return
res
;
}
public
:
NumArray
(
vector
<
int
>&
nums
)
{
// nlogn
bit
=
vector
<
int
>
(
nums
。
size
()
+
1
,
0
);
// O(nlogn)建樹
data
=
vector
<
int
>
(
nums
。
size
(),
0
);
for
(
int
i
=
0
;
i
<
nums
。
size
();
i
++
)
update
(
i
,
nums
[
i
]);
// // O(n)建樹
// data = nums;
// for(int i = 1; i <= nums。size(); i++){
// bit[i] += nums[i-1];
// int j = i + lowBit(i);
// if(j <= nums。size()) bit[j] += bit[i];
// }
}
void
update
(
int
i
,
int
val
)
{
//logn
int
diff
=
val
-
data
[
i
];
data
[
i
]
=
val
;
i
+=
1
;
// bit中下標從1開始
while
(
i
<=
data
。
size
()){
bit
[
i
]
+=
diff
;
i
+=
lowBit
(
i
);
}
}
int
sumRange
(
int
i
,
int
j
)
{
//logn
return
sum
(
j
+
1
)
-
sum
(
i
);
}
};
2。4 解法三、線段樹
2。4。1 線段樹(Segment Tree)
和樹狀陣列一樣,利用線段樹也可以以對數複雜度進行區間查詢和元素更新。樹狀陣列是利用索引的二進位制表示來劃分子序列,而線段樹是不斷進行二分來劃分子序列,如下圖所示
可知線段樹是一棵二叉樹,樹中的每個結點儲存著一個區間範圍以及對應區間和(也可以儲存區間最大值、最小值、異或值等等,只要滿足結合律就可以了),其左右子結點分別儲存該結點區間拆分為兩半之後各自區間的資訊。 我們就可以自頂向下不斷查詢得到某個區間和,也可以自底向上地不斷更新某個元素對應的葉結點到根結點的路徑上的結點值來進行元素更新,由於樹高是O(logn),所以區間查詢和元素修改複雜度都是O(logn)。
我們可以採用自頂向下不斷遞迴直到葉結點的建樹方式;也可以採用自底向上迭代的建樹方式。前者就是普通的線段樹,需要不斷遞迴壓棧;而後者就是常用的zkw線段樹(出自張昆瑋的論文《統計的力量》,zkw即其名字拼音首字母),本文采用的是後者。
我們用一個數組st(segment tree)來存放線段樹且下標從1開始,先來看看當資料陣列nums的大小n剛好是2的冪時的情況,此時最簡單,因為構建出來的剛好是一個完美二叉樹,如下圖所示
上圖所示的是一棵完美二叉樹,葉子數(即資料陣列nums的大小4)剛好是2的冪,此時st的大小是2n(算上了無用的st[0]),而且可立即得到第一個葉子結點nums[0]存放在st[4],第二個葉子結點nums[1]存放在st[5],以此類推。
即葉子結點對應nums下標剛好與st相差n,其中n是nums的大小。
所以我們很好寫出當nums大小剛好是2的冪時的建樹程式碼:
// nums的大小n是2的冪時, 建立線段樹, 複雜度O(n)
void
buildST
(
vector
<
int
>&
nums
)
{
for
(
int
i
=
n
;
i
<
n
*
2
;
++
i
)
st
[
i
]
=
nums
[
i
-
n
];
// 葉子
for
(
int
i
=
n
-
1
;
i
>
0
;
——
i
)
st
[
i
]
=
st
[
i
<<
1
]
+
st
[(
i
<<
1
)
|
1
];
// st[i << 1]和st[(i << 1)|1]分別表示左右孩子
}
然後自底向上修改某個元素則很簡單:
// nums的大小n是2的冪時, 更新元素, 複雜度O(logn)
void
update
(
int
i
,
int
val
)
{
i
+=
n
;
// 轉成到st的下標
st
[
i
]
=
val
;
// 更新葉子
while
(
i
>
0
)
{
// st[i^1]表示st[i]的兄弟
st
[
i
>>
1
]
=
st
[
i
]
+
st
[
i
^
1
];
i
>>=
1
;
}
}
如何進行自頂向下的區間查詢
sumRange(i, j)
呢? 1。 若st[i]是右子結點,說明結果應包含它但不包含它的父親,那麼將結果加上st[i]並使i增加1,最後將i除以2進入下一迴圈; 2。 若st[i]是左子結點,它跟其右兄弟在要求的區間裡,則此時直接將i除以2(即直接進入其父親結點)進入下一迴圈即可;
對j的處理同理:若j是左子結點,那麼需要加上st[j]並使j減去1最後將j除以2進入下一迴圈;若j是右子結點,直接將j除以2進入下一迴圈即可。可以透過判斷i的奇偶性來判斷st[i]是左子結點還是右子結點。
// nums的大小n是2的冪時, 區間查詢, 複雜度O(logn)
int
sumRange
(
int
i
,
int
j
)
{
i
+=
n
;
j
+=
n
;
// 轉成到st的下標
int
res
=
0
;
for
(;
i
<=
j
;
i
>>=
1
,
j
>>=
1
){
if
((
i
&
1
))
res
+=
st
[
i
++
];
// st[i]是右子結點
if
(
!
(
j
&
1
))
res
+=
st
[
j
——
];
// st[j]是左子結點
}
return
res
;
}
上述分析的前提是nums的大小n是2的冪,但是對於一般情況下又是如何呢?
簡單的方法是在nums的結尾補0,直到其長度正好為2的冪。最壞情況下,我們需要大小約為4n的st而不是2n(其實問題也不大,很多人就是這麼做的)。例如,若n=17,則我們應該補15個0,開闢大小為64的st。
其實,
前面的程式碼對於任意大小的nums都是正確的!
我們依然只需要開闢一個2n大小的陣列st就可以按照上述程式碼進行建樹、查詢、更新操作。
二叉樹中: 度為0的結點數 = 度為2的結點數 + 1,又因為st中沒有度為1的結點,所以只需開闢2n大小的st陣列(算上了無用的st[0])。
例如若n=6,那麼建立的st樹如下,可以驗證一下前面查詢、更新操作的程式碼都是正確的。
另外,線段樹還支援區間更新,此時可以開闢一個大小為n的懶標記(lazy tag)陣列來提高效率,由於此題並沒有涉及,所以不詳述。
2。4。2 程式碼
將前面的程式碼總結一下即可得到完整程式碼:
class
NumArray
{
private
:
int
n
;
// nums。size()
vector
<
int
>
st
;
// segment tree
// int leftChild(int r){return r << 1;}
// int rightChild(int r){return r << 1 | 1;} // r << 1 + 1
void
buildST
(
vector
<
int
>&
nums
)
{
for
(
int
i
=
n
;
i
<
n
*
2
;
++
i
)
st
[
i
]
=
nums
[
i
-
n
];
// 葉子
for
(
int
i
=
n
-
1
;
i
>
0
;
——
i
)
st
[
i
]
=
st
[
i
<<
1
]
+
st
[(
i
<<
1
)
|
1
];
// st[i << 1]和st[(i << 1)|1]分別表示左右孩子
}
public
:
NumArray
(
vector
<
int
>&
nums
)
{
n
=
nums
。
size
();
st
=
vector
<
int
>
(
2
*
n
,
0
);
buildST
(
nums
);
}
void
update
(
int
i
,
int
val
)
{
i
+=
n
;
// 轉成到st的下標
st
[
i
]
=
val
;
while
(
i
>
0
)
{
// st[i^1]表示st[i]的兄弟
st
[
i
>>
1
]
=
st
[
i
]
+
st
[
i
^
1
];
i
>>=
1
;
}
}
int
sumRange
(
int
i
,
int
j
)
{
i
+=
n
;
j
+=
n
;
// 轉成到st的下標
int
res
=
0
;
for
(;
i
<=
j
;
i
>>=
1
,
j
>>=
1
){
if
((
i
&
1
))
res
+=
st
[
i
++
];
// st[i]是右子結點
if
(
!
(
j
&
1
))
res
+=
st
[
j
——
];
// st[j]是左子結點
}
return
res
;
}
};
2。5 總結
本文講的三個解法都是對陣列先劃分成子序列,只是劃分的方式不同。2。2節的折中思路就是無腦將陣列均分成大小為sqrt(n)的塊,區間查詢和單點修改的複雜度都是O(sqrt(n));樹狀陣列則採用精心設計的劃分方式,區間查詢和單點修改的複雜度都是O(logn);而線段樹則是不斷進行二分,區間查詢和單點修改的複雜度也都是O(logn)。
樹狀陣列和線段樹的思想很類似,不過也有不同之處,具體區別和聯絡如下:
樹狀陣列邏輯上是一棵普通的樹,而線段樹邏輯上是一顆完全二叉樹;
兩者時間複雜度級別相同, 但是樹狀陣列的常數明顯優於線段樹而且程式碼實現簡單;
線段樹的空間複雜度在常數上為樹狀陣列的兩倍;
一般來講,凡是可以使用樹狀陣列解決的問題, 使用線段樹也可以解決, 但是線段樹能夠解決的問題樹狀陣列未必能夠解決(例如求區間最大/小值);
參考
樹狀陣列(Binary Indexed Tree),看這一篇就夠了 - CSDN
樹狀陣列詳解 - 部落格園
Efficient and easy segment trees - Codeforces
關於線段樹(Segment tree)和樹狀陣列(BIT)的區別 - 知乎
更多我的LeetCode中文題解,可前往GitHub檢視:
https://
github。com/ShusenTang/L
eetCode