線段樹(Segment Tree)是常用的維護區間資訊的資料結構,它可以在 O(logn) 的時間複雜度下實現單點修改、區間修改、區間查詢(區間求和、區間最大值或區間最小值)等操作,常用來解決 RMQ 問題。
RMQ(Range Minimum/Maximum Query) 問題是指:對於長度為 n 的數列 A,回答若干詢問 RMQ(A, i, j) 其中 i, j <= n,返回數列 A 中下標在 i, j 裡的最小(大)值。也就是說:RMQ問題是指求區間最值的問題。通常該型別題目的解法有遞迴分治、動態規劃、線段樹和單調棧/單調佇列。
這篇內容斷斷續續寫了兩週,隨著練習對線段樹的理解不斷深入,慢慢地學習下來也不覺得它有多麼困難,更多的體會還是熟能生巧,雖然它起初看上去確實程式碼量大一些,但是我覺得只要大家放平心態,循序漸進的掌握下文中的三部分,也沒什麼難的。
1. 線段樹
線段樹會將每個長度不為 1 的區間劃分成左右兩個區間來遞迴求解,透過合併左右兩區間的資訊來求得當前區間的資訊。
比如,我們將一個大小為 5 的陣列 nums = {10, 11, 12, 13, 14} 轉換成線段樹,並規定線段樹的根節點編號為 1。用陣列 tree[] 來儲存線段樹的節點,tree[i] 表示線段樹上編號為 i 的節點,圖示如下:
圖示中每個節點展示了區間和以及區間範圍,tree[i] 左子樹節點為 tree[2i],右子樹節點為 tree[2i + 1]。如果 tree[i] 記錄的區間為 [a, b] 的話,那麼左子樹節點記錄的區間為 [a, mid],右子樹節點記錄的區間為 [mid + 1, b],其中 mid = (a + b) / 2。
現在我們已經對線段樹有了基本的認識,接下來我們看看區間查詢和單點修改的程式碼實現。
區間查詢和單點修改線段樹
首先,我們定義線段樹的節點:
/**
* 定義線段樹節點
*/
class Node {
/**
* 區間和 或 區間最大/最小值
*/
int val;
int left;
int right;
public Node(int left, int right) {
this.left = left;
this.right = right;
}
}
注意其中的 val 欄位儲存的是區間的和。定義完樹的節點,我們來看一下建樹的邏輯,注意程式碼中的註釋,我們為線段樹分配的節點陣列大小為原陣列大小的 4 倍,這是考慮到陣列轉換成滿二叉樹的最壞情況。
public SegmentTree(int[] nums) {
this.nums = nums;
tree = new Node[nums.length * 4];
// 建樹,注意表示區間時使用的是從 1 開始的索引值
build(1, 1, nums.length);
}
/**
* 建樹
*
* @param pos 當前節點編號
* @param left 當前節點區間下界
* @param right 當前節點區間上界
*/
private void build(int pos, int left, int right) {
// 建立節點
tree[pos] = new Node(left, right);
// 遞迴結束條件
if (left == right) {
// 賦值
tree[pos].val = nums[left - 1];
return;
}
// 如果沒有到根節點,則繼續遞迴
int mid = left + right >> 1;
build(pos << 1, left, mid);
build(pos << 1 | 1, mid + 1, right);
// 當前節點的值是左子樹和右子樹節點的和
pushUp(pos);
}
/**
* 用於向上回溯時修改父節點的值
*/
private void pushUp(int pos) {
tree[pos].val = tree[pos << 1].val + tree[pos << 1 | 1].val;
}
我們在建樹時,表示區間並不是從索引 0 開始,而是從索引 1 開始,這樣才能保證在計算左子樹節點索引時為 2i,右子樹節點索引為 2i + 1。
build()
方法執行時,我們會先在對應的位置上建立節點而不進行賦值,只有在遞迴到葉子節點時才賦值,此時區間大小為 1,節點值即為當前區間的值。之後非葉子節點值都是透過pushUp()
方法回溯加和當前節點的兩個子節點值得出來的。
接下來我們看修改區間中的值,線段樹對值的更新方法,關注其中的註釋:
/**
* 修改單節點的值
*
* @param pos 當前節點編號
* @param numPos 需要修改的區間中值的位置
* @param val 修改後的值
*/
private void update(int pos, int numPos, int val) {
// 找到該數值所線上段樹中的葉子節點
if (tree[pos].left == numPos && tree[pos].right == numPos) {
tree[pos].val = val;
return;
}
// 如果不是當前節點那麼需要判斷是去左或右去找
int mid = tree[pos].left + tree[pos].right >> 1;
if (numPos <= mid) {
update(pos << 1, numPos, val);
} else {
update(pos << 1 | 1, numPos, val);
}
// 葉子節點的值修改完了,需要回溯更新所有相關父節點的值
pushUp(pos);
}
修改方法比較簡單,當葉子節點值更新完畢時,我們仍然需要呼叫pushUp()
方法對所有相關父節點值進行更新。
接下來我們看查詢對應區間和的方法:
/**
* 查詢對應區間的值
*
* @param pos 當前節點
* @param left 要查詢的區間的下界
* @param right 要查詢的區間的上界
*/
private int query(int pos, int left, int right) {
// 如果我們要查詢的區間把當前節點區間全部包含起來
if (left <= tree[pos].left && tree[pos].right <= right) {
return tree[pos].val;
}
int res = 0;
int mid = tree[pos].left + tree[pos].right >> 1;
// 根據區間範圍去左右節點分別查詢求和
if (left <= mid) {
res += query(pos << 1, left, right);
}
if (right > mid) {
res += query(pos << 1 | 1, left, right);
}
return res;
}
該方法也比較簡單,需要判斷區間範圍是否需要對向左子節點和右子節點的分別查詢計算。
現在表示區間和的線段樹已經講解完畢了,為了方便大家學習和看程式碼,我把全量的程式碼在這裡貼出來:
public class SegmentTree {
/**
* 定義線段樹節點
*/
static class Node {
/**
* 區間和 或 區間最大/最小值
*/
int val;
int left;
int right;
public Node(int left, int right) {
this.left = left;
this.right = right;
}
}
Node[] tree;
int[] nums;
public SegmentTree(int[] nums) {
this.nums = nums;
tree = new Node[nums.length * 4];
// 建樹,注意表示區間時使用的是從 1 開始的索引值
build(1, 1, nums.length);
}
/**
* 建樹
*
* @param pos 當前節點編號
* @param left 當前節點區間下界
* @param right 當前節點區間上界
*/
private void build(int pos, int left, int right) {
// 建立節點
tree[pos] = new Node(left, right);
// 遞迴結束條件
if (left == right) {
// 賦值
tree[pos].val = nums[left - 1];
return;
}
// 如果沒有到根節點,則繼續遞迴
int mid = left + right >> 1;
build(pos << 1, left, mid);
build(pos << 1 | 1, mid + 1, right);
// 當前節點的值是左子樹和右子樹節點的和
pushUp(pos);
}
/**
* 修改單節點的值
*
* @param pos 當前節點編號
* @param numPos 需要修改的區間中值的位置
* @param val 修改後的值
*/
private void update(int pos, int numPos, int val) {
// 找到該數值所線上段樹種的葉子節點
if (tree[pos].left == numPos && tree[pos].right == numPos) {
tree[pos].val = val;
return;
}
// 如果不是當前節點那麼需要判斷是去左或右去找
int mid = tree[pos].left + tree[pos].right >> 1;
if (numPos <= mid) {
update(pos << 1, numPos, val);
} else {
update(pos << 1 | 1, numPos, val);
}
// 葉子節點的值修改完了,需要回溯更新所有相關父節點的值
pushUp(pos);
}
/**
* 用於向上回溯時修改父節點的值
*/
private void pushUp(int pos) {
tree[pos].val = tree[pos << 1].val + tree[pos << 1 | 1].val;
}
/**
* 查詢對應區間的值
*
* @param pos 當前節點
* @param left 要查詢的區間的下界
* @param right 要查詢的區間的上界
*/
private int query(int pos, int left, int right) {
// 如果我們要查詢的區間把當前節點區間全部包含起來
if (left <= tree[pos].left && tree[pos].right <= right) {
return tree[pos].val;
}
int res = 0;
int mid = tree[pos].left + tree[pos].right >> 1;
// 根據區間範圍去左右節點分別查詢求和
if (left <= mid) {
res += query(pos << 1, left, right);
}
if (right > mid) {
res += query(pos << 1 | 1, left, right);
}
return res;
}
}
如果要建立表示區間最大值或最小值的線段樹,建樹的邏輯不變,只需要將pushUp()方法和query()方法修改成計算最大值或最小值的邏輯即可。
2. 線段樹的區間修改與懶惰標記
如果我們不僅對單點進行修改,也對區間進行修改,那麼在區間修改時就需要將當前區間值及包含當前區間的子區間值都修改一遍,這樣所產生的開銷是沒辦法接受的,因此在這裡我們會使用一種懶惰標記的方法來幫助我們避免這種即時開銷。
簡單來說,懶惰標記是透過延遲對節點資訊的更改,從而減少可能不必要的操作次數。每次執行修改時,我們透過打標記的方法表明該節點對應的區間在某一次操作中被更改,但不更新該節點的子節點的資訊。實質性的修改則在下一次“即將訪問(update 或 query)”到帶有懶惰標記節點的子節點時才進行。
我們透過在節點類中新增 add 欄位記錄懶惰標記,它表示的是該區間的子區間值需要“變化的大小”(一定好好好的理解),並透過 pushDown 方法“累加”到當前區間的兩個子節點區間值中。
只要不訪問到當前區間的子區間,那麼子區間值始終都不會變化,相當於子區間值的變化量被當前節點透過 add 欄位“持有”
pushDown 方法區別於我們上文中提到的 pushUp 方法,前者是將當前節點值累計的懶惰標記值同步到子節點中,而後者是完成子節點修改後,回溯修改當前子節點的父節點值,我們能夠根據 Down 和 Up 來更好的理解這兩個方法的作用方向和修改範圍。
下面我們一起來看看過程和具體的程式碼,節點類如下,增加 add 欄位:
static class Node {
int left;
int right;
int val;
int add;
public Node(int left, int right) {
this.left = left;
this.right = right;
}
}
區間修改
建樹的流程與我們上述的一致,就不在這裡贅述了,我們主要關注區間修改的過程,還是以如下初始的線段樹為例,此時各個節點的 add 均為 0:
接下來我們修改區間 [3, 5] 且區間內每個值變化量為 1,過程如下:
先遍歷節點 1,發現 [3, 5] 區間不能將 [1, 5] 區間完全包含,不進行修改,繼續遍歷節點 2。節點 2 依然沒有被區間 [3, 5] 包含,需要繼續遍歷節點 5,發現該節點被區間完全包含,進行修改並新增懶惰標記值,如下圖所示:
完成這一步驟後需要向上回溯修改 tree[2] 節點的值:
現在 [3, 5] 區間中 3 已經完成修改,還有 4, 5 沒有被修改,我們需要在右子樹中繼續遞迴查詢,發現 tree[3] 中區間被我們要修改的區間 [3, 5] 完全包含,那麼需要將這個節點進行修改並懶惰標記,如下,注意這裡雖然 tree[3] 節點有兩個子節點,但是因為我們沒有訪問到它的子節點所以無需同步 add 值到各個子節點中:
同樣,完成這一步驟也需要向上回溯修改父節點的值:
到現在我們的區間修改就已經完成了,根據這個過程程式碼示例如下:
/**
* 修改區間的值
*
* @param pos 當前節點編號
* @param left 要修改區間的下界
* @param right 要修改區間的上界
* @param val 區間內每個值的變化量
*/
public void update(int pos, int left, int right, int val) {
// 如果該區間被要修改的區間包圍的話,那麼需要將該區間所有的值都修改
if (left <= tree[pos].left && tree[pos].right <= right) {
tree[pos].val += (tree[pos].right - tree[pos].left + 1) * val;
// 懶惰標記
tree[pos].add += val;
return;
}
// 該區間沒有被包圍的話,需要修改節點的資訊
pushDown(pos);
int mid = tree[pos].left + tree[pos].right >> 1;
// 如果下界在 mid 左邊,那麼左子樹需要修改
if (left <= mid) {
update(pos << 1, left, right, val);
}
// 如果上界在 mid 右邊,那麼右子樹也需要修改
if (right > mid) {
update(pos << 1 | 1, left, right, val);
}
// 修改完成後向上回溯修改父節點的值
pushUp(pos);
}
private void pushDown(int pos) {
// 根節點 和 懶惰標記為 0 的情況不需要再向下遍歷
if (tree[pos].left != tree[pos].right && tree[pos].add != 0) {
int add = tree[pos].add;
// 計算累加變化量
tree[pos << 1].val += add * (tree[pos << 1].right - tree[pos << 1].left + 1);
tree[pos << 1 | 1].val += add * (tree[pos << 1 | 1].right - tree[pos << 1 | 1].left + 1);
// 子節點懶惰標記累加
tree[pos << 1].add += add;
tree[pos << 1 | 1].add += add;
// 懶惰標記清 0
tree[pos].add = 0;
}
}
private void pushUp(int pos) {
tree[pos].val = tree[pos << 1].val + tree[pos << 1 | 1].val;
}
區間查詢
tree[3] 節點是有懶惰標記 1 的,如果我們此時查詢區間 [5, 5] 的值,就需要在遞迴經過 tree[3] 節點時,進行 pushDown 懶惰標記計算,將 tree[6] 和 tree[7] 的節點值進行修改,結果如下:
最終我們會獲取到結果值為 15,區間查詢過程的示例程式碼如下:
public int query(int pos, int left, int right) {
if (left <= tree[pos].left && tree[pos].right <= right) {
// 當前區間被包圍
return tree[pos].val;
}
// 懶惰標記需要下傳修改子節點的值
pushDown(pos);
int res = 0;
int mid = tree[pos].left + tree[pos].right >> 1;
if (left <= mid) {
res += query(pos << 1, left, right);
}
if (right > mid) {
res += query(pos << 1 | 1, left, right);
}
return res;
}
同樣,為了方便大家學習,我把全量程式碼也列出來,我認為學習線段樹的區間修改比較重要的點是理解 add 欄位表示的含義和 pushDown 方法的作用時機,而且需要注意只有線段樹中的某個區間被我們要修改的區間全部包含時(update 和 query 方法的條件判斷),才進行值修改並懶惰標記,否則該區間值只在 pushUp 方法回溯時被修改。
public class SegmentTree2 {
static class Node {
int left;
int right;
int val;
int add;
public Node(int left, int right) {
this.left = left;
this.right = right;
}
}
Node[] tree;
int[] nums;
public SegmentTree2(int[] nums) {
this.tree = new Node[nums.length * 4];
this.nums = nums;
build(1, 1, nums.length);
}
private void build(int pos, int left, int right) {
tree[pos] = new Node(left, right);
// 遞迴結束條件
if (left == right) {
tree[pos].val = nums[left - 1];
return;
}
int mid = left + right >> 1;
build(pos << 1, left, mid);
build(pos << 1 | 1, mid + 1, right);
// 回溯修改父節點的值
pushUp(pos);
}
/**
* 修改區間的值
*
* @param pos 當前節點編號
* @param left 要修改區間的下界
* @param right 要修改區間的上界
* @param val 區間內每個值的變化量
*/
public void update(int pos, int left, int right, int val) {
// 如果該區間被要修改的區間包圍的話,那麼需要將該區間所有的值都修改
if (left <= tree[pos].left && tree[pos].right <= right) {
tree[pos].val += (tree[pos].right - tree[pos].left + 1) * val;
// 懶惰標記
tree[pos].add += val;
return;
}
// 該區間沒有被包圍的話,需要修改節點的資訊
pushDown(pos);
int mid = tree[pos].left + tree[pos].right >> 1;
// 如果下界在 mid 左邊,那麼左子樹需要修改
if (left <= mid) {
update(pos << 1, left, right, val);
}
// 如果上界在 mid 右邊,那麼右子樹也需要修改
if (right > mid) {
update(pos << 1 | 1, left, right, val);
}
// 修改完成後向上回溯修改父節點的值
pushUp(pos);
}
public int query(int pos, int left, int right) {
if (left <= tree[pos].left && tree[pos].right <= right) {
// 當前區間被包圍
return tree[pos].val;
}
// 懶惰標記需要下傳修改子節點的值
pushDown(pos);
int res = 0;
int mid = tree[pos].left + tree[pos].right >> 1;
if (left <= mid) {
res += query(pos << 1, left, right);
}
if (right > mid) {
res += query(pos << 1 | 1, left, right);
}
return res;
}
private void pushDown(int pos) {
// 根節點 和 懶惰標記為 0 的情況不需要再向下遍歷
if (tree[pos].left != tree[pos].right && tree[pos].add != 0) {
int add = tree[pos].add;
// 計算累加變化量
tree[pos << 1].val += add * (tree[pos << 1].right - tree[pos << 1].left + 1);
tree[pos << 1 | 1].val += add * (tree[pos << 1 | 1].right - tree[pos << 1 | 1].left + 1);
// 子節點懶惰標記
tree[pos << 1].add += add;
tree[pos << 1 | 1].add += add;
// 懶惰標記清 0
tree[pos].add = 0;
}
}
private void pushUp(int pos) {
tree[pos].val = tree[pos << 1].val + tree[pos << 1 | 1].val;
}
}
3. 線段樹動態開點
線段樹的動態開點其實不難理解,它與我們上述直接建立好線段樹所有節點不同,動態開點的線段樹在最初只建立一個根節點代表整個區間,其他節點只在需要的時候被建立,節省出了空間。當然,我們因此也不能再使用pos << 1
和pos << 1 | 1
來尋找當前節點的左右子節點,取而代之的是在節點中使用 left 和 right 記錄左右子節點在 tree[] 中的位置,這一點需要注意:
static class Node {
// left 和 right 不再表示區間範圍而是表示左右子節點在 tree 中的索引位置
int left, right;
int val;
int add;
}
我們以區間 [1, 5] 為例,建立區間 [5, 5] 為 14 的過程圖示如下:
我們可以發現,會先建立預設的根節點 tree[1],之後建立出上圖中 tree[2] 和 tree[3] 節點,而此時並沒有找到區間 [5, 5],那麼需要繼續建立上圖中的 tree[4] 和 tree[5] 節點(與直接建立出所有節點不同,如果是直接建立好所有節點的話它們的位置應該在 tree[6] 和 tree[7]),現在 tree[5] 節點表示的區間符合我們要找的條件,可以進行賦值和 pushUp 操作了,與直接建立出所有節點相比,動態開點少建立了 4 個節點,也就是圖中標紅的四個節點我們是沒有建立的。
由於每次操作都可能建立並訪問全新的一系列節點,因此 m 次單點操作後節點的空間複雜度是 O(mlogn),如果我們採用線段樹動態開點解題的話,空間要開的儘可能大,Java 在 128M 可以開到 5e6 個節點以上。
結合圖示大家應該能理解動態開點的過程了(不明白就自己畫一遍),下面我們看下具體的程式碼:
/**
* 修改區間的值
*
* @param pos 當前節點的索引值
* @param left 當前線段樹節點表示的範圍下界
* @param right 當前線段樹節點表示的範圍上界
* @param l 要修改的區間下界
* @param r 要修改的區間上界
* @param val 區間值變化的大小
*/
public void update(int pos, int left, int right, int l, int r, int val) {
// 當前區間被要修改的區間全部包含
if (l <= left && right <= r) {
tree[pos].val += (right - left + 1) * val;
tree[pos].add += val;
return;
}
lazyCreate(pos);
pushDown(pos, right - left + 1);
int mid = left + right >> 1;
if (l <= mid) {
update(tree[pos].left, left, mid, l, r, val);
}
if (r > mid) {
update(tree[pos].right, mid + 1, right, l, r, val);
}
pushUp(pos);
}
// 為該位置建立節點
private void lazyCreate(int pos) {
if (tree[pos] == null) {
tree[pos] = new Node();
}
// 建立左子樹節點
if (tree[pos].left == 0) {
tree[pos].left = ++count;
tree[tree[pos].left] = new Node();
}
// 建立右子樹節點
if (tree[pos].right == 0) {
tree[pos].right = ++count;
tree[tree[pos].right] = new Node();
}
}
private void pushDown(int pos, int len) {
if (tree[pos].left != 0 && tree[pos].right != 0 && tree[pos].add != 0) {
// 計算左右子樹的值
tree[tree[pos].left].val += (len - len / 2) * tree[pos].add;
tree[tree[pos].right].val += len / 2 * tree[pos].add;
// 子節點懶惰標記
tree[tree[pos].left].add += tree[pos].add;
tree[tree[pos].right].add += tree[pos].add;
tree[pos].add = 0;
}
}
private void pushUp(int pos) {
tree[pos].val = tree[tree[pos].left].val + tree[tree[pos].right].val;
}
整體的邏輯並不難,新增的 lazyCreate 方法是動態開點的邏輯,需要注意的是執行區間更新時我們方法的引數中多了 left 和 right 表示當前節點區間範圍的引數,因為我們現在的節點中只儲存了左右子節點的位置,而沒有區間資訊,所以我們需要在引數中攜帶才行,否則我們沒有辦法判斷當前區間和要找的區間是否匹配。
我還是將全量程式碼放在下面,方便大家學習:
public class SegmentTree3 {
static class Node {
// left 和 right 不再表示區間範圍而是表示左右子節點在 tree 中的索引位置
int left, right;
int val;
int add;
}
// 記錄當前節點數
int count;
Node[] tree;
public SegmentTree3() {
count = 1;
tree = new Node[(int) 5e6];
tree[count] = new Node();
}
public int query(int pos, int left, int right, int l, int r) {
if (l <= left && right <= r) {
return tree[pos].val;
}
lazyCreate(pos);
pushDown(pos, right - left + 1);
int res = 0;
int mid = left + right >> 1;
if (l <= mid) {
res += query(tree[pos].left, left, mid, l, r);
}
if (r > mid) {
res += query(tree[pos].right, mid + 1, right, l, r);
}
return res;
}
/**
* 修改區間的值
*
* @param pos 當前節點的索引值
* @param left 當前線段樹節點表示的範圍下界
* @param right 當前線段樹節點表示的範圍上界
* @param l 要修改的區間下界
* @param r 要修改的區間上界
* @param val 區間值變化的大小
*/
public void update(int pos, int left, int right, int l, int r, int val) {
// 當前區間被要修改的區間全部包含
if (l <= left && right <= r) {
tree[pos].val += (right - left + 1) * val;
tree[pos].add += val;
return;
}
lazyCreate(pos);
pushDown(pos, right - left + 1);
int mid = left + right >> 1;
if (l <= mid) {
update(tree[pos].left, left, mid, l, r, val);
}
if (r > mid) {
update(tree[pos].right, mid + 1, right, l, r, val);
}
pushUp(pos);
}
// 為該位置建立節點
private void lazyCreate(int pos) {
if (tree[pos] == null) {
tree[pos] = new Node();
}
// 建立左子樹節點
if (tree[pos].left == 0) {
tree[pos].left = ++count;
tree[tree[pos].left] = new Node();
}
// 建立右子樹節點
if (tree[pos].right == 0) {
tree[pos].right = ++count;
tree[tree[pos].right] = new Node();
}
}
private void pushDown(int pos, int len) {
if (tree[pos].left != 0 && tree[pos].right != 0 && tree[pos].add != 0) {
// 計算左右子樹的值
tree[tree[pos].left].val += (len - len / 2) * tree[pos].add;
tree[tree[pos].right].val += len / 2 * tree[pos].add;
// 子節點懶惰標記
tree[tree[pos].left].add += tree[pos].add;
tree[tree[pos].right].add += tree[pos].add;
tree[pos].add = 0;
}
}
private void pushUp(int pos) {
tree[pos].val = tree[tree[pos].left].val + tree[tree[pos].right].val;
}
}
巨人的肩膀
作者:京東物流 王奕龍
來源:京東雲開發者社群 自猿其說 Tech 轉載請註明來源