Java數(shù)據(jù)結構之線段樹詳解

更新時間：2022年01月23日 10:44:00 作者：strongmore

線段樹是一種二叉搜索樹，與區(qū)間樹相似，它將一個區(qū)間劃分成一些單元區(qū)間，每個單元區(qū)間對應線段樹中的一個葉結點。本文將介紹線段樹的Java實現(xiàn)代碼，需要的可以參考一下

介紹

線段樹(又名區(qū)間樹)也是一種二叉樹，每個節(jié)點的值等于左右孩子節(jié)點值的和，線段樹示例圖如下

以求和為例，根節(jié)點表示區(qū)間0-5的和，左孩子表示區(qū)間0-2的和，右孩子表示區(qū)間3-5的和，依次類推。

代碼實現(xiàn)

/**
 * 使用數(shù)組實現(xiàn)線段樹
 */
public class SegmentTree<E> {

  private Node[] data;
  private int size;

  private Merger<E> merger;

  public SegmentTree(E[] source, Merger<E> merger) {
    this.merger = merger;
    this.size = source.length;
    this.data = new Node[size * 4];
    buildTree(0, source, 0, size - 1);
  }

  public E search(int queryLeft, int queryRight) {
    if (queryLeft < 0 || queryLeft > size || queryRight < 0 || queryRight > size
        || queryLeft > queryRight) {
      throw new IllegalArgumentException("index is illegal");
    }
    return search(0, queryLeft, queryRight);
  }

  /**
   * 查詢區(qū)間queryLeft-queryRight的值
   */
  private E search(int treeIndex, int queryLeft, int queryRight) {
    Node treeNode = data[treeIndex];
    int left = treeNode.left;
    int right = treeNode.right;
    if (left == queryLeft && right == queryRight) {
      return elementData(treeIndex);
    }
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    int middle = left + ((right - left) >> 1);
    if (queryLeft > middle) {
      return search(rightTreeIndex, queryLeft, queryRight);
    } else if (queryRight <= middle) {
      return search(leftTreeIndex, queryLeft, queryRight);
    }
    E leftEle = search(leftTreeIndex, queryLeft, middle);
    E rightEle = search(rightTreeIndex, middle + 1, queryRight);
    return merger.merge(leftEle, rightEle);
  }

  public void update(int index, E e) {
    update(0, index, e);
  }

  /**
   * 更新索引為index的值為e
   */
  private void update(int treeIndex, int index, E e) {
    Node treeNode = data[treeIndex];
    int left = treeNode.left;
    int right = treeNode.right;
    if (left == right) {
      treeNode.data = e;
      return;
    }
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    int middle = left + ((right - left) >> 1);
    if (index > middle) {
      update(rightTreeIndex, index, e);
    } else {
      update(leftTreeIndex, index, e);
    }
    treeNode.data = merger.merge(elementData(leftTreeIndex), elementData(rightTreeIndex));
  }

  private void buildTree(int treeIndex, E[] source, int left, int right) {
    if (left == right) {
      data[treeIndex] = new Node<>(source[left], left, right);
      return;
    }
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    int middle = left + ((right - left) >> 1);
    buildTree(leftTreeIndex, source, left, middle);
    buildTree(rightTreeIndex, source, middle + 1, right);
    E treeData = merger.merge(elementData(leftTreeIndex), elementData(rightTreeIndex));
    data[treeIndex] = new Node<>(treeData, left, right);
  }

  @Override
  public String toString() {
    return Arrays.toString(data);
  }

  private E elementData(int index) {
    return (E) data[index].data;
  }

  private int leftChild(int index) {
    return index * 2 + 1;
  }

  private int rightChild(int index) {
    return index * 2 + 2;
  }

  private static class Node<E> {

    E data;
    int left;
    int right;

    Node(E data, int left, int right) {
      this.data = data;
      this.left = left;
      this.right = right;
    }

    @Override
    public String toString() {
      return String.valueOf(data);
    }
  }

  public interface Merger<E> {

    E merge(E e1, E e2);
  }
}

我們以LeetCode上的一個問題來分析線段樹的構建，查詢和更新，LeetCode307問題如下：

給定一個整數(shù)數(shù)組，查詢索引區(qū)間[i,j]的元素的總和。

線段樹構建

private void buildTree(int treeIndex, E[] source, int left, int right) {
    if (left == right) {
      data[treeIndex] = new Node<>(source[left], left, right);
      return;
    }
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    int middle = left + ((right - left) >> 1);
    buildTree(leftTreeIndex, source, left, middle);
    buildTree(rightTreeIndex, source, middle + 1, right);
    E treeData = merger.merge(elementData(leftTreeIndex), elementData(rightTreeIndex));
    data[treeIndex] = new Node<>(treeData, left, right);
  }

測試代碼

public class Main {

  public static void main(String[] args) {
    Integer[] nums = {-2, 0, 3, -5, 2, -1};
    SegmentTree<Integer> segmentTree = new SegmentTree<>(nums, Integer::sum);
    System.out.println(segmentTree);
  }

}

最后構造出的線段樹如下，前面為元素值，括號中為包含的區(qū)間。

遞歸構造過程為

當左指針和右指針相等時，表示為葉子節(jié)點
將左孩子和右孩子值相加，構造當前節(jié)點，依次類推

區(qū)間查詢

  /**
   * 查詢區(qū)間queryLeft-queryRight的值
   */
  private E search(int treeIndex, int queryLeft, int queryRight) {
    Node treeNode = data[treeIndex];
    int left = treeNode.left;
    int right = treeNode.right;
    if (left == queryLeft && right == queryRight) {
      return elementData(treeIndex);
    }
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    int middle = left + ((right - left) >> 1);
    if (queryLeft > middle) {
      return search(rightTreeIndex, queryLeft, queryRight);
    } else if (queryRight <= middle) {
      return search(leftTreeIndex, queryLeft, queryRight);
    }
    E leftEle = search(leftTreeIndex, queryLeft, middle);
    E rightEle = search(rightTreeIndex, middle + 1, queryRight);
    return merger.merge(leftEle, rightEle);
  }

查詢區(qū)間2-5的和

public class Main {

  public static void main(String[] args) {
    Integer[] nums = {-2, 0, 3, -5, 2, -1};
    SegmentTree<Integer> segmentTree = new SegmentTree<>(nums, Integer::sum);
    System.out.println(segmentTree);
    System.out.println(segmentTree.search(2, 5)); // -1
  }

}

查詢過程為

待查詢的區(qū)間和當前節(jié)點的區(qū)間相等，返回當前節(jié)點值
待查詢左區(qū)間大于中間區(qū)間值，查詢右孩子
待查詢右區(qū)間小于中間區(qū)間值，查詢左孩子
待查詢左區(qū)間在左孩子，右區(qū)間在右孩子，兩邊查詢結果相加

更新

  /**
   * 更新索引為index的值為e
   */
  private void update(int treeIndex, int index, E e) {
    Node treeNode = data[treeIndex];
    int left = treeNode.left;
    int right = treeNode.right;
    if (left == right) {
      treeNode.data = e;
      return;
    }
    int leftTreeIndex = leftChild(treeIndex);
    int rightTreeIndex = rightChild(treeIndex);
    int middle = left + ((right - left) >> 1);
    if (index > middle) {
      update(rightTreeIndex, index, e);
    } else {
      update(leftTreeIndex, index, e);
    }
    treeNode.data = merger.merge(elementData(leftTreeIndex), elementData(rightTreeIndex));
  }

更新只影響元素值，不影響元素區(qū)間。

更新其實和構建的邏輯類似，找到待更新的實際索引，依次更新父節(jié)點的值。