[程式] Segment Tree (線段樹)

Segment Tree 是個用空間來交換時間的資料結構，特別適合用在搜尋次數極為龐大的時候。其所需的額外空間約為既有資料的兩倍 (空間複雜度是 O(N)!)，時間複雜度則為 O(log(N))。

資料結構的部分就請直接參考一開頭的連結，或者網路上搜尋都有很多，這邊就不再介紹。以下僅就實作部分略作說明：

首先考量到要支援 C++ 的既有型態以及使用者自定義的型態，所以一樣是採用 template 來處理型態問題。當然這種方法的老問題就是沒辦法確定該型態有哪些運算子可用，因此這邊就學 C++ 一些會用到比較運算的容器，預設會用到 operator< (程式中 CHOICE 設為 2 的範例)。

不過進一步考量到自定義型態的 operator< 可能希望預設用在其他演算法 (ex: sort，或是 heap 之類的)，而 segment tree 這邊要的其實不單是 "比較" 這樣的概念，而是希望根據左子樹與右子樹的數值做完某種 "運算" 後得到一個結果 (例如計算總合，參考 CHOICE 為 6 的範例)，這時候應該要提供一個方法讓使用者可以自定義這種行為。因此藉由第二個 template parameter 來讓使用者可以自行設定。(另一個範例：CHOICE 為 1)

當然，藉由這種方法 SegmentTree 本身使用上的彈性就大多了，不過這樣還是有一項不足之處：使用者其實也可以為自定義型態提供 operator() 使其變為具有 function 特性的 functor，這樣就無須另外寫一個 struct 或 class (參考 CHOICE 為 3 的範例)。

綜合以上，我們就可以寫出一個相當具有彈性的資料結構，不過這邊有個小麻煩：當使用者自定義型態同時多載 (overload) 了 operator< 跟 operator() 該怎麼辦呢？這邊我的選擇是採用 operator() 而非 operator<，理由相當簡單：operator() 比起 operator< 更加貼近原本的意思。因此，我們必須提供一個方法去檢查使用者是否有提供 operator()，如果有的話則優先選用 operator()；反之，則選用預設的 operator<。

要達成上述的要求，這邊利用的是 C++ template 的一個概念：SFINAE (substitution failure is not an error)，可以參考程式碼的第 25 ~ 82 行，其中 struct has_function 就是利用 SFINAE 的主軸，根據給定的型態是否有 operator() 可以讓 has_function 的成員變數 value 分別為 true 或 false。而後續的 struct DefaultSegTreeOp 就可以利用 has_function 來判斷究竟有沒有 operator() 可以用。SFINAE 的概念也可以參考我以前寫過的這篇，有另一個針對 primitive type 跟 container type 做區別的用法。

以下程式碼 (目前只有 query 跟 update，update 尚不支援一個區間，僅可更新單一節點)：

	/*
	* =====================================================================================
	*
	* Filename: SegmentTree.cpp
	*
	* Description: Implementation of the data structure: segment tree
	*
	* Version: 1.0
	* Created: 2018/02/18 (yyyy/mm/dd)
	* Revision: none
	* Compiler: g++14
	*
	* Author: lionking
	* Organization: None
	*
	* =====================================================================================
	*/
	#include <cstdio>
	#include <vector>
	#include <iostream>
	#define CHOICE 0 // enter a number to choose an example
	// test whether the given type has overloaded operator() or not
	// concept: SFINAE
	template <class T>
	struct has_function
	{
	typedef char true_type;
	typedef long false_type;
	// template not available if there's no overloaded operator< in U's scope
	// Note: operator() must be non-static member function
	template <class U>
	static true_type test(typeof (&U::operator()));
	// fallback
	template <class U>
	static false_type test(...);
	// tests which overload of test is chosen for T
	static const bool value = sizeof(test<T>(0)) == sizeof(true_type);
	};
	// pre-declaration
	template <bool b>
	struct DefaultSegTreeOpImpl;
	// specialization: situation when the given type exist overloaded operator()
	template <>
	struct DefaultSegTreeOpImpl<true>
	{
	template <typename T>
	T operator() (const T& lhs, const T& rhs) const
	{
	T op;
	return op(lhs, rhs);
	}
	};
	// specialization: situation when the given type do not have overloaded operator()
	template <>
	struct DefaultSegTreeOpImpl<false>
	{
	template <typename T>
	T operator() (const T& lhs, const T& rhs) const { return std::min(lhs, rhs); }
	};
	struct DefaultSegTreeOp
	{
	template <typename T>
	T operator() (const T& lhs, const T& rhs) const
	{
	DefaultSegTreeOpImpl<has_function<T>::value> op;
	return op(lhs, rhs);
	}
	};
	/*
	* Tree node for constructing segment tree
	*/
	template <typename T>
	struct SegTreeNode
	{
	size_t first, last; // specify this tree node can cover the original data within [first, last)
	T data;
	SegTreeNode(size_t l=0, size_t r=0) : first(l), last(r) {};
	SegTreeNode(const T& new_data, size_t l=0, size_t r=0) : first(l), last(r), data(new_data) {};
	void set(const T& new_data, size_t l, size_t r) { data = new_data; first = l; last = r; }
	};
	/********************************************************************************************************
	* Segment tree
	* Given an array of data, this class will construct a segment tree for query in O(log |N|),
	* where N is the number of data in the given array.
	*
	* template parameter:
	* 1. T: type of the given data
	* 2. Operator is a binary function that accepts two arguments with type T and returns a result in type T
	* default: std::min
	*
	* Note:
	* When T is a user-defined structure, it must overload at least one of the following operators:
	* 1. operator<
	* 2. operator()
	* When both of the above operators are overloaded, operator() will be chosen for execution.
	********************************************************************************************************/
	template < typename T, typename Operator=DefaultSegTreeOp >
	class SegmentTree
	{
	public:
	typedef T value_type;
	typedef Operator op_type;
	// parameters:
	// first, last:
	// Random-access iterators to the initial and final positions of the sequence to be constructed in segment tree.
	// The range used is [first,last), which contains all the elements between first and last,
	// including the element pointed by first but not the element pointed by last.
	template <typename RandomAccessIterator>
	void construct(RandomAccessIterator first, RandomAccessIterator last);
	// query data within range [first, last)
	// Note: first and last are indices in original data set,
	// and index starts from 0
	inline T query(const size_t first, const size_t last) const;
	// update data on the specified index
	// Note: parameter index is the index in original data set and starts from 0
	inline void update(const size_t index, const T& data);
	private:
	std::vector< SegTreeNode<T> > tree;
	Operator compare;
	static constexpr size_t tree_root = 0;
	inline size_t getlc(const size_t node) const { return ((node << 1) + 1); }
	inline size_t getrc(const size_t node) const { return ((node + 1) << 1); }
	// sub-function for construct (segment tree construction)
	template <typename RandomAccessIterator>
	void construct(const RandomAccessIterator& first, size_t curr, size_t start, size_t end);
	// query data within range [first, last)
	// curr: currntly queried tree node
	T query(const size_t curr, const size_t first, const size_t last) const;
	// update segment tree
	// curr: currntly updated tree node
	void update(const size_t index, const T& data, const size_t curr);
	};
	template <typename T, typename Operator>
	template <typename RandomAccessIterator>
	void SegmentTree<T, Operator>::construct(const RandomAccessIterator& first, size_t curr, size_t start, size_t end)
	{
	if (start == (end - 1)) {
	tree[curr].set(*(first + start), start, end);
	}
	else if (start < end) {
	const auto middle = (start + end) >> 1;
	const auto lchild = getlc(curr);
	const auto rchild = getrc(curr);
	construct(first, lchild, start, middle); // recursive to left child
	construct(first, rchild, middle, end); // recursive to right child
	tree[curr].set(compare(tree[lchild].data, tree[rchild].data), start, end);
	}
	else { return; }
	}
	template <typename T, typename Operator>
	template <typename RandomAccessIterator>
	void SegmentTree<T, Operator>::construct(RandomAccessIterator first, RandomAccessIterator last)
	{
	const size_t size = std::distance(first, last);
	// Minimum required leaf nodes is the minimum x s.t. 2^x >= number of elements
	size_t tree_size = 1;
	for (; (tree_size < size); tree_size <<= 1);
	// Then, minimum required tree node is 2^(x+1)
	tree.resize(tree_size << 1);
	// segment tree construction
	construct(first, tree_root, 0, size);
	}
	template <typename T, typename Operator>
	T SegmentTree<T, Operator>::query(const size_t curr, const size_t first, const size_t last) const
	{
	const auto &curr_node = tree.at(curr);
	const auto lchild = getlc(curr);
	const auto rchild = getrc(curr);
	if ((first >= curr_node.last) || (last < curr_node.first)) {
	// current node cannot cover query range.
	return T();
	}
	else if ((first <= curr_node.first) && (curr_node.last <= last)) {
	// query range has fully cover current node
	// return data stored in current node;
	return curr_node.data;
	}
	else if (last <= tree.at(lchild).last) {
	// left child can fully cover query range
	return query(lchild, first, last);
	}
	else if (tree.at(rchild).first <= first) {
	// right chile can fully cover query range
	return query(rchild, first, last);
	}
	else {
	const auto& lhs = query(lchild, first, std::min(last, tree.at(lchild).last));
	const auto& rhs = query(rchild, std::max(first, tree.at(rchild).first), last);
	return compare(lhs, rhs);
	}
	}
	template <typename T, typename Operator>
	inline T SegmentTree<T, Operator>::query(const size_t first, const size_t last) const
	{
	return query(tree_root, first, last);
	}
	template <typename T, typename Operator>
	void SegmentTree<T, Operator>::update(const size_t index, const T& data, const size_t curr)
	{
	auto& curr_node = tree.at(curr);
	if ((index < curr_node.first) || (curr_node.last <= index)) {
	// tree nodes that should be updated is not within the range corresponding to currently processed node
	return;
	}
	else if ((index == curr_node.first) && ((curr_node.last - curr_node.first) == 1)) {
	// target leaf node, update data directly
	curr_node.data = data;
	}
	else {
	const auto lchild = getlc(curr);
	const auto rchild = getrc(curr);
	update(index, data, lchild); // update data in left child
	update(index, data, rchild); // update data in right child
	// merge: update data in current node
	curr_node.data = compare(tree.at(lchild).data, tree.at(rchild).data);
	}
	}
	template <typename T, typename Operator>
	inline void SegmentTree<T, Operator>::update(const size_t index, const T& data)
	{
	update(index, data, tree_root);
	}
	/****************************************************
	* The followings are examples of using SegmentTree *
	****************************************************/
	struct SegTreeOp1
	{
	template <typename T>
	T operator() (const T& lhs, const T& rhs) const { return std::max(lhs, rhs); }
	};
	struct SegTreeOp2
	{
	template <typename T>
	T operator() (const T& lhs, const T& rhs) const { return (lhs + rhs); }
	};
	struct STNode1
	{
	int value;
	STNode1(int v = 0) : value(v) {}
	bool operator< (const STNode1& rhs) const { return value < rhs.value; }
	};
	struct STNode2
	{
	int value;
	STNode2(int v = 0) : value(v) {}
	STNode2 operator() (const STNode2& lhs, const STNode2& rhs) const { return std::max(lhs.value, rhs.value); }
	};
	struct STNode3
	{
	int value;
	STNode3(int v = 0) : value(v) {}
	bool operator< (const STNode3& rhs) const { return value < rhs.value; }
	STNode3 operator() (const STNode3& lhs, const STNode3& rhs) const { return std::max(lhs.value, rhs.value); }
	};
	struct STNode4
	{
	int value;
	STNode4(int v = 0) : value(v) {}
	};
	bool operator< (const STNode4& lhs, const STNode4& rhs) { return lhs.value < rhs.value; }
	int main()
	{
	std::cout << "choice: " << CHOICE << '\n';
	#if CHOICE == 0
	// The simplest example:
	// 1. Use primitive type
	// 2. No extra comparator is specified for segment tree
	const int data[] = {1, 5, 2, 3, 4};
	SegmentTree<int> st;
	st.construct(data, data + 5);
	std::cout << st.query(1, 5) << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5) << '\n';
	#elif CHOICE == 1
	// 1. Use primitive type
	// 2. Extra comparator is specified for segment tree
	const int data[] = {1, 5, 2, 3, 4};
	SegmentTree<int, SegTreeOp1> st;
	st.construct(data, data + 5);
	std::cout << st.query(1, 5) << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5) << '\n';
	#elif CHOICE == 2
	// The simplest example of using user-defined data structure.
	// At least one of operator< or operator() must be provided.
	SegmentTree<STNode1> st;
	const STNode1 data[] = {1, 5, 2, 3, 4};
	st.construct(data, data + 5);
	std::cout << st.query(1, 5).value << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5).value << '\n';
	#elif CHOICE == 3
	// The simplest example of using user-defined data structure.
	// At least one of operator< or operator() must be provided.
	SegmentTree<STNode2> st;
	const STNode2 data[] = {1, 5, 2, 3, 4};
	st.construct(data, data + 5);
	std::cout << st.query(1, 5).value << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5).value << '\n';
	#elif CHOICE == 4
	// when both operator< and operator() are provided, SegmentTree will choose operator() for comparison
	SegmentTree<STNode3> st;
	const STNode3 data[] = {1, 5, 2, 3, 4};
	st.construct(data, data + 5);
	std::cout << st.query(1, 5).value << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5).value << '\n';
	#elif CHOICE == 5
	// Different from choice 2, the overloaded operator< is non-member function.
	SegmentTree<STNode4> st;
	const STNode4 data[] = {1, 5, 2, 3, 4};
	st.construct(data, data + 5);
	std::cout << st.query(1, 5).value << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5).value << '\n';
	#elif CHOICE == 6
	// 1. Use primitive type
	// 2. Extra operator is specified for segment tree
	const int data[] = {1, 5, 2, 3, 4};
	SegmentTree<int, SegTreeOp2> st;
	st.construct(data, data + 5);
	std::cout << st.query(1, 5) << '\n';
	std::cout << st.query(2, 4) << '\n';
	st.update(1, -1);
	std::cout << st.query(1, 5) << '\n';
	std::cout << st.query(2, 4) << '\n';
	#else
	std::cout << "choice " << CHOICE << " is not available.\n";
	#endif
	return 0;
	}

view raw SegmentTree.cpp hosted with ❤ by GitHub