數據結構之線段樹

什麼是線段樹

線段樹(Segment Tree)也叫區間樹,其本質上是一種二分搜索樹,不一樣點在於線段樹中每一個節點再也不是存放單純的元素,而是存放了一個能夠表示區間的值,一般是該區間合併後的值。而且每一個區間會被平均分爲2個子區間,做爲它的左右子節點。好比說根節點存放了區間 [1,10],那麼就會被分爲區間 [1,5] 做爲左子節點,區間 [6,10] 做爲右子節點。java

例如,咱們能夠將這樣一個數組所表示的區間構形成線段樹:
數據結構之線段樹數組

而且指定區間合併規則爲區間內的元素求和,那麼構造出來的線段樹表示以下:
數據結構之線段樹bash

  • 從這顆線段樹能夠看到,因爲具備二分搜索樹的特性,咱們能夠快速地在線段樹中找到一個區間。注意,這裏是指按區間查找,而不是按元素值查找。因此線段樹相對難理解的地方就在於每一個節點既有區間的概念又有一個元素值。

爲何要使用線段樹

關於線段樹的一個經典問題就是:區間染色。假設有一面牆,長度爲 n,每次選擇一段兒牆進行染色。在 m 次操做後,咱們能夠在 [i, j] 區間內看見多少中顏色?數據結構

對於這個問題,咱們可使用一個數組來實現:
數據結構之線段樹app

對於染色操做(更新區間)咱們能夠遍歷數組找到目標區間進行染色,時間複雜度是 $O(n)$。對於查詢操做(查詢區間)也是遍歷數組便可,一樣時間複雜度爲 $O(n)$。顯然用線性結構來解決這類問題的時間複雜度要更高一些,此時線段樹就派上用場了,由於樹形結構的時間複雜度一般在 $O(logn)$。ide

除此以外,線段樹的另外一個經典問題就是:區間查詢。查詢一個區間 [i, j] 的最大值和最小值,或者區間數字之和。例如,在實際業務中很常見的基於區間的統計查詢:2017年註冊用戶中消費最高的用戶?消費最少的用戶?學習時間最長的用戶?某個太空區間中天體總量?學習

對於靜態區間數據(區間內的數據不會發生變化)來講,是比較好解決的,但以上所提到的問題都是動態的區間數據(區間內的數據在不斷的變化),此時線段樹就是一個比較好的選擇。測試

經過以上的介紹,咱們能總結出線段樹的兩個核心操做:ui

  • 區間更新:更新區間中一個元素或者一個區間的值
  • 區間查詢:查詢一個區間 [i, j] 的最大值、最小值,或者區間數字之和

線段樹基礎表示

線段樹雖然不像堆那樣是一棵徹底二叉樹,但線段樹因爲其特性知足平衡二叉樹(左右子樹高度相差不超過1),因此依然可使用數組進行表示。咱們能夠將其看作是一顆滿二叉樹,空節點就當作葉子節點便可。以下示例:
數據結構之線段樹this

既然能夠用數組來表示一棵線段樹,那麼若是區間有 n 個元素,此時應該建立多大容量的數組來構建一顆線段樹呢?對於這個問題,咱們先來看如何求一棵滿二叉樹的節點:假設這棵樹有 h 層,那麼這棵樹就一共有 $2^h-1$ 個節點(大約是 $2^h$)。對於最後一層($h - 1$ 層)來講,就有 $2^{(h-1)}$ 個節點。所以,最後一層的節點數大體等於前面全部層節點之和。

瞭解瞭如何求滿二叉樹的節點數量後,回到以前的問題,若是區間有 n 個元素,此時應該開多大空間的數組?咱們能夠分紅兩種狀況:

  • 若是 $n = 2^k$,那麼只須要開闢 $2n$ 的數組空間
  • 若是 $n = 2^k + 1$,那麼就須要開闢 $4n$ 的數組空間

一般來講,咱們的線段樹不考慮添加元素,即區間固定(區間內的數據能夠是不固定的),那麼使用 $4n$ 的靜態空間便可。這也是廣泛構造線段樹時,使用的一個通用值。除非對內存有嚴格要求,不然通常開闢 $4n$ 的數組空間便可。並且對於內存有要求的狀況下,通常也不會採用數組來表示,此時鏈式結會是更優的選擇。

接下來,咱們就實現一下線段樹的基礎結構代碼:

package tree;

/**
 * 線段樹 - 基於數組的表示實現
 *
 * @author 01
 * @date 2021-01-27
 **/
public class SegmentTree<E> {

    /**
     * 保存原始數組,即須要被構形成線段樹的區間
     */
    private E[] data;

    /**
     * 線段樹的數組表示
     */
    private E[] tree;

    public SegmentTree(E[] arr) {
        this.data = (E[]) new Object[arr.length];
        System.arraycopy(arr, 0, this.data, 0, arr.length);

        // 開闢 4n 的數組空間用於構造線段樹
        this.tree = (E[]) new Object[4 * arr.length];
    }

    public int getSize() {
        return data.length;
    }

    public E get(int index) {
        if (index < 0 || index >= data.length) {
            throw new IllegalArgumentException("Index is illegal");
        }

        return data[index];
    }

    /**
     * 返回徹底二叉樹的數組表示中,一個索引所表示的元素的左子節點的索引
     */
    private int leftChild(int index) {
        return 2 * index + 1;
    }

    /**
     * 返回徹底二叉樹的數組表示中,一個索引所表示的元素的右子節點的索引
     */
    private int rightChild(int index) {
        return 2 * index + 2;
    }
}

建立線段樹

在本小節中,咱們來根據以前實現的基礎代碼,完成建立線段樹邏輯的編寫。須要說明一下的是,在本例中,線段樹每一個節點所存儲的元素是區間合併後的值。具體的實現代碼以下:

/**
 * 用戶自定義的區間合併邏輯
 */
private final Merger<E> merger;

public SegmentTree(E[] arr, Merger<E> merger) {
    this.merger = merger;
    this.data = (E[]) new Object[arr.length];
    System.arraycopy(arr, 0, this.data, 0, arr.length);

    // 開闢 4n 的數組空間用於構建線段樹
    this.tree = (E[]) new Object[4 * arr.length];
    // 構建線段樹,傳入根節點索引,以及區間的左右端點
    buildSegmentTree(0, 0, data.length - 1);
}

/**
 * 在treeIndex的位置建立表示區間[left...right]的線段樹
 */
private void buildSegmentTree(int treeIndex, int left, int right) {
    // 區間中只有一個元素,表明遞歸到底了
    if (left == right) {
        tree[treeIndex] = data[left];
        return;
    }

    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    // 計算中間點,須要避免整型溢出
    int mid = left + (right - left) / 2;
    // 構建左子樹
    buildSegmentTree(leftTreeIndex, left, mid);
    // 構建右子樹
    buildSegmentTree(rightTreeIndex, mid + 1, right);

    // 對於兩個區間的合併規則是與業務相關的,因此要調用用戶自定義的邏輯來完成
    tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
}

/**
 * 遍歷打印樹中節點中值信息。
 *
 * @return String
 */
@Override
public String toString() {
    StringBuilder res = new StringBuilder();
    res.append('[');
    for (int i = 0; i < tree.length; i++) {
        if (tree[i] != null) {
            res.append(tree[i]);
        } else {
            res.append("null");
        }

        if (i != tree.length - 1) {
            res.append(", ");
        }
    }
    res.append(']');

    return res.toString();
}
  • 在線段樹中根節點存儲的數據,實際就是左右兩個子節點數據的合併(遞歸便可),而具體如何合併是由業務決定的。例如,能夠是求和,也能夠是求最大值或最小值。另外,這裏沒有經過一個對象來表示節點中的左右區間,而是經過方法參數的形式表示了這個區間,數組中只存儲區間合併後的值。

用戶傳入的 Merger 是一個接口,其定義以下:

package tree;

/**
 * 合併器接口
 *
 * @author 01
 * @date 2021-01-27
 **/
public interface Merger<E> {

    /**
     * 用戶自定義的區間合併邏輯
     *
     * @param a 區間a
     * @param b 區間b
     * @return 合併後的結果
     */
    E merge(E a, E b);
}

最後,咱們來編寫一個簡單的測試用例進行一下測試:

package tree;

/**
 * 測試SegmentTree
 *
 * @author 01
 */
public class SegmentTreeTests {

    public static void main(String[] args) {
        Integer[] nums = {-2, 0, 3, -5, 2, -1};
        SegmentTree<Integer> segTree = new SegmentTree<>(
                nums, Integer::sum // 對兩個區間中的值進行求和
        );
        System.out.println(segTree);
    }
}

輸出結果以下:

[-3, 1, -4, -2, 3, -3, -1, -2, 0, null, null, -5, 2, null, null, null, null, null, null, null, null, null, null, null]
  • 能夠看到,線段樹的根節點是 -3 ,由於對整個數組的求和結果就是 -3 。左子節點爲 1,由於 -2 + 0 + 3 = 1。右子節點爲 -4,同理,由於 -5 + 2 + -1 = -4,其他以此類推。結果符合預期,證實咱們實現的線段樹沒有問題。

線段樹中的區間查詢

例如,咱們要對以下這棵線段樹查詢 [2, 5] 這個區間:
數據結構之線段樹

因爲咱們以前傳入的 Merger 實現的是求和邏輯,那麼這至關於查詢2 ~ 5區間全部元素的和。從根節點開始往下,咱們知道分割位置,左節點查詢 [2, 3],右節點查詢 [4, 5],找到兩個節點以後合併就能夠了。

具體的實現代碼以下:

/**
 * 查詢區間[queryLeft, queryRight]的值,如[2, 5]
 */
public E query(int queryLeft, int queryRight) {
    if (queryLeft < 0 || queryLeft >= data.length ||
            queryRight < 0 || queryRight >= data.length ||
            queryLeft > queryRight) {
        throw new IllegalArgumentException("Index is illegal");
    }

    return query(0, 0,
            data.length - 1, queryLeft, queryRight);
}

/**
 * 在以treeIndex爲根的線段樹中[left...right]的範圍裏,搜索區間[queryLeft...queryRight]的值
 */
private E query(int treeIndex, int left, int right,
                int queryLeft, int queryRight) {
    // 找到了目標區間
    if (left == queryLeft && right == queryRight) {
        return tree[treeIndex];
    }

    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    // 計算中間點,須要避免整型溢出
    int mid = left + (right - left) / 2;

    if (queryLeft >= mid + 1) {
        // 目標區間不在左子樹中,查找右子樹
        return query(rightTreeIndex, mid + 1, right, queryLeft, queryRight);
    } else if (queryRight <= mid) {
        // 目標區間不在右子樹中,查找左子樹
        return query(leftTreeIndex, left, mid, queryLeft, queryRight);
    }

    // 目標區間一部分在右子樹中,一部分在左子樹中,則兩個子樹都須要找
    E leftResult = query(leftTreeIndex, left, mid, queryLeft, mid);
    E rightResult = query(rightTreeIndex, mid + 1, right, mid + 1, queryRight);

    // 找到目標區間的值,將其合併後返回
    return merger.merge(leftResult, rightResult);
}

進行一個簡單的測試:

public static void main(String[] args) {
    Integer[] nums = {-2, 0, 3, -5, 2, -1};
    SegmentTree<Integer> segTree = new SegmentTree<>(
            nums, Integer::sum // 對兩個區間中的值進行求和
    );

    System.out.println(segTree.query(0,2));
    System.out.println(segTree.query(2,5));
    System.out.println(segTree.query(0,5));
}

輸出結果以下:

1
-1
-3

線段樹中的更新操做

咱們使用線段樹來解決區間相關的問題,主要是針對區間內的數據是動態變化的狀況,若是是靜態區間通常不須要用到線段樹。因此在本小節,咱們就來實現線段樹中的更新操做。

實際上線段樹中的更新操做,本質上是在二分查找。由於根據線段樹的特性,待更新的目標節點確定是一個葉子節點,咱們只須要找到這個葉子節點並進行更新便可。咱們查找待更新節點的依據是數組的索引,而數組的索引是從 0 ~ n 有序的,因此在一個有序的區間中查找某個特定的值,妥妥的就是二分查找了。

知道了咱們在更新線段樹中某個節點時,要找的這個待更新節點是一個葉子節點,而且找到這個葉子節點的過程本質上是一個二分查找,那麼這個思路就很清晰了。

首先,將找到葉子節點的條件做爲遞歸的退出條件。而後計算中間點,並將線段樹數組劃分爲 [left...mid][mid+1...right] 兩個區間。接着判斷要找的數組索引落在哪一個區間,就繼續往哪一個區間遞歸查找。最後,將區間的值進行合併。如此一來,就完成了目標節點的更新操做。

具體的實現代碼以下:

/**
 * 將index位置的值,更新爲e
 */
public void set(int index, E e) {
    if (index < 0 || index >= data.length) {
        throw new IllegalArgumentException("Index is illegal");
    }

    data[index] = e;
    set(0, 0, data.length - 1, index, e);
}

/**
 * 在以treeIndex爲根的線段樹中更新index的值爲e
 */
private void set(int treeIndex, int left, int right, int index, E e) {
    // 找到了葉子節點
    if (left == right) {
        // 進行更新
        tree[treeIndex] = e;
        return;
    }

    int mid = left + (right - left) / 2;
    // 將線段樹數組劃分爲[left...mid]和[mid+1...right]兩個區間
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    if (index >= mid + 1) {
        // index在右子樹
        set(rightTreeIndex, mid + 1, right, index, e);
    } else {
        // index在左子樹
        set(leftTreeIndex, left, mid, index, e);
    }

    tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
}

Leetcode上線段樹相關的問題

在本文的最後,咱們來使用本身實現的線段樹解決一個Leetcode上的307號問題:

該問題的主要需求是更新數組下標對應的值,以及查詢數組中某個區間內的元素總和。像這種對區間內數據有更新需求的,會使得區間內數據動態變化的,就很適合使用線段樹來解決。具體的實現代碼以下:

package tree.solution;

import tree.SegmentTree;

/**
 * Leetcode 307. Range Sum Query - Mutable
 * https://leetcode.com/problems/range-sum-query-mutable/description/
 */
class NumArray {

    private SegmentTree<Integer> segTree;

    public NumArray(int[] nums) {
        if (nums.length != 0) {
            Integer[] data = new Integer[nums.length];
            for (int i = 0; i < nums.length; i++) {
                data[i] = nums[i];
            }
            segTree = new SegmentTree<>(data, Integer::sum);
        }
    }

    public void update(int i, int val) {
        if (segTree == null) {
            throw new IllegalArgumentException("Error");
        }
        segTree.set(i, val);
    }

    public int sumRange(int i, int j) {
        if (segTree == null) {
            throw new IllegalArgumentException("Error");
        }

        return segTree.query(i, j);
    }
}
相關文章
相關標籤/搜索