数据结构与算法之线段树

(一) 定义

线段树是一棵平衡的二叉搜索树. 何为平衡? 线段树的左右两个子树的高度差的绝对值不超过1. 线段树的每个结点存储的内容包括两部分:

• 区间或者是线段: [区间左端点, 区间右端点] (以参数的形式表示 或者 节点类的成员属性表示)
• 区间结点存储的元素: 根据业务需求的来决定

基于平衡二叉树的特性(左右子树的高度差的绝对值不超过1), 我们可以用NULL补全不存在的区间, 最终补全为一个满二叉树, 此时将区间节点按顺序一层一层的码放出来(层序输出), 因此我们可以使用使用数组的方式表示一颗完全二叉树. (跟区间存在n个元素, 数组最大需要4n的空间大小存储)

由上图可知:

• 每个节点的左孩子区间范围为[l, mid], 右孩子为[mid+1, r]. 其中 l: 左端点, r: 右端点, mid = ( l + r ) / 2
• 对于结点i, 左孩子结点为 2*i + 1, 右孩子为 2*i + 2.

(二) 自定义线段树

1.根据数组创建线段树

public class SegmentTree<E> {

	/**
	 * 源数据数组
	 */
	private E[] data;

	/**
	 * 以数组的形式表现线段树
	 */
	private E[] tree;
	
	/**
	 * 融合器接口: 线段树的区间结点存储的元素根据融合器接口的merger()方法来决定
	 */
	private Merger<E> merger;

	/**
	 * 构造函数: 将传来的数组 构建成 线段树数组
	 * 
	 * @param arr
	 */
	@SuppressWarnings("unchecked")
	public SegmentTree(E[] arr, Merger<E> merger) {
		// 融合器接口初始化 
		this.merger = merger;
		
		// 源数据数组初始化
		data = (E[]) new Object[arr.length];
		for (int i = 0; i < arr.length; i++) {
			data[i] = arr[i];
		}

		// 线段树数组初始化的大小为 数组大小的4倍.
		tree = (E[]) new Object[4 * arr.length];
		// 构建线段树数组
		buildSegmentTree(0, 0, arr.length - 1);
	}

	/**
	 * 在treeIndex的位置创建表示区间[l...r]的线段树
	 * 
	 * @param treeIndex	线段树的根结点索引
	 * @param l	区间的左端点
	 * @param r	区间的右端点
	 */
	private void buildSegmentTree(int treeIndex, int l, int r) {
		// 递归的终止条件: 区间的左端点 等于 区间的右端点
		if (l == r) {
			tree[treeIndex] = data[l];
			return;
		}
		
		// 当前结点的左孩子结点索引
		int leftTreeIndex = leftChild(treeIndex);
		// 当前结点的右孩子结点索引
		int rightTreeIndex = rightChild(treeIndex);
		
		// 中间值
		int mid = l + (r - l) / 2;
		
		// 先创建以左孩子结点索引处区间[最左端点...中间值]的线段树
		buildSegmentTree(leftTreeIndex, l, mid);
		// 然后创建以右孩子结点索引处区间[中间值+1...最右端点]的线段树
        buildSegmentTree(rightTreeIndex, mid + 1, r);
		
		// 最后区间要维护的数据 根据Merger接口的merge(e1, e2)方法来决定
		tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
	}
	
	/**
	 * 返回完全二叉树的数组表示中, 一个索引所表示的元素的左孩子结点的索引
	 *  
	 * @param index
	 * @return
	 */
	private int leftChild(int index) {
		return index * 2 + 1;
	}
	
	/**
	 * 返回完全二叉树的数组表示中, 一个索引所表示的元素的右孩子结点的索引
	 * 
	 * @param index
	 * @return
	 */
	private int rightChild(int index) {
		return index * 2 + 2;
	}
	
	@Override
	public String toString() {
		StringBuilder res = new StringBuilder();
		res.append('[');
		for (int i = 0; i < tree.length; i++) {
			res.append(tree[i]);
			if (i != tree.length - 1) {
				res.append(", ");
			}
		}
		res.append("]");
		return res.toString();
	}
}

测试

public static void main(String[] args) {
	Integer[] nums = {1, 2, 3, 4, 5, 6};
	// lambda表达式 (a, b) -> a + b 表示Merger接口的merge(a, b), 求两数之和
	SegmentTree<Integer> segmentTree = new SegmentTree<>(nums, (a, b) -> a + b);
	System.out.println(segmentTree);
}

2.线段树的区间查询操作

存在三种情况:

• 线段树的左孩子区间包含搜索区间(mid >= 查询区间的右端点 ), 递归进左孩子区间
• 线段树的右孩子区间包含搜索区间(mid <= 查询区间的左端点 ), 递归进右孩子区间
• 搜索区间同时存在线段树的左右孩子区间, 获取 [查询区间的左端点, mid] 和 [mid, 查询区间的右端点] 的值, 进行操作.

/**
 * 返回区间[queryL, queryR]的值
 * 
 * @param queryL
 * @param queryR
 * @return
 */
public E query(int queryL, int queryR) {
	 if (queryL < 0 || queryL >= data.length || queryR < 0 || queryR >= data.length || queryR < queryL) {
		 throw new IllegalArgumentException("Index is illegel.");
	 }
	 return query(0, 0, data.length - 1, queryL, queryR);
}

/**
 * 在以treeIndex为根的线段树中[r...l]的范围里, 搜索区间[queryL...queryR]的值
 * 
 * @param treeIndex
 * @param l
 * @param r
 * @param queryL
 * @param queryR
 * @return
 */
private E query(int treeIndex, int l, int r, int queryL, int queryR) {
	// 递归终止条件: 搜索区间的左端点 等于 以treeIndex为根的线段树的左端点  且 搜索区间的右端点 等于 以treeIndex为根的线段树的右端点
	if (l == queryL && r == queryR) {
		return tree[treeIndex];
	}
	
	int leftTreeIndex = leftChild(treeIndex);
	int rightTreeIndex = rightChild(treeIndex);
	int mid = l + (r - l) / 2; // 3 
	
	if (queryR <= mid) {
		// 第一种情况: 以treeIndex为根的线段树的左子树区间  包含 搜索区间
		return query(leftTreeIndex, l, mid, queryL, queryR);
	} else if (queryL >= mid + 1) {
		// 第二种情况: 以treeIndex为根的线段树的右子树区间  包含 搜索区间
		return query(rightTreeIndex, mid + 1, r, queryL, queryR);
	} else {
		// 第三种情况: 搜索区间 包含在  以treeIndex为根的线段树的左右子树区间
		E leftResult = query(leftTreeIndex, l, mid, queryL, mid);
		E rightResult = query(rightTreeIndex, mid + 1, r, mid + 1, queryR);
		return this.merger.merge(leftResult, rightResult);
	}
}

3.线段树的单个更新操作

在线段树中更新index索引对应的值, 且维护其所有父节点的的值.

/**
 * 将index位置的值, 更新为e
 * 
 * @param index
 * @param 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
 * 
 * @param treeIndex
 * @param l
 * @param r
 * @param index
 * @param e
 */
private void set(int treeIndex, int l, int r, int index, E e) {
	// 递归终止的条件: 搜索区间的左端点 等于 搜索区间的右端点, 找到index索引在线段树的位置
	if (l == r) {
		tree[treeIndex] = e;
		return;
	}
	
	int leftTreeIndex = leftChild(treeIndex);
	int rightTreeIndex = rightChild(treeIndex);
	int mid = l + (r - l) / 2;
	
	if (index <= mid) {
		// 以treeIndex为根的线段树的左子树区间  包含 index
		set(leftTreeIndex, l, mid, index, e);
	} else { // index >= mid + 1
		// 以treeIndex为根的线段树的右子树区间  包含 index
		set(rightTreeIndex, mid + 1, r, index, e);
	}
	
	// 对修改元素结点的父节点重新赋值merge
	tree[treeIndex] = merger.merge(tree[leftTreeIndex], tree[rightTreeIndex]);
}

(三) 时间复杂度分析

函数	时间复杂度	分析
query(l, r)	O(h) => O(logn)	自定义线段树是一颗满二叉树, 查询操作只需在树的高度上递归
set(index, e)	O(h) => O(logn)	自定义线段树是一颗满二叉树, 更新操作只需在树的高度上递归