实现

  1. class SegmentTree:
  2.     def __init__(self, data: list):
  3.         # 线段树
  4.         self.tree = [None for i in range(len(data) * 4)]
  5.         self.data = [i for i in data]
  7.     def get_size(self):
  8.         return len(self.data)
  10.     def get(self, index: int):
  11.         assert 0 <= index < len(self.data), "index error"
  12.         return self.data[index]
  14.     @staticmethod
  15.     def get_left_child(index):
  16.         """ 左孩子节点下标"""
  17.         return 2 * index + 1
  19.     @staticmethod
  20.     def get_right_child(index):
  21.         """右孩子节点下标"""
  22.         return 2 * index + 2
  24.     def build_segment_tree(self):
  25.         self.__build_segment_tree(0, 0, len(self.data) - 1)
  27.     # 在tree_index的位置创建表示区间[l,r]的线段树
  28.     def __build_segment_tree(self, tree_index, l: int, r: int):
  29.         if l == r:
  30.             self.tree[tree_index] = self.data[l]
  31.             return
  32.         # 左下标
  33.         left_index = self.get_left_child(tree_index)
  34.         # 右下标
  35.         right_index = self.get_right_child(tree_index)
  36.         mid = l + (r - l) // 2
  37.         self.__build_segment_tree(left_index, l, mid)
  38.         self.__build_segment_tree(right_index, mid + 1, r)
  39.         self.tree[tree_index] = self.tree[left_index] + self.tree[right_index]
  41.     def query(self, query_l, query_r):
  42.         return self.__query(0, 0, len(self.data) - 1, query_l, query_r)
  44.     # 在以tree_index 为根的线段树中[l..r]的范围内，搜索区间[query_l,query_r]的值
  45.     def __query(self, tree_index, l, r, query_l, query_r):
  46.         if l == query_l and r == query_r:
  47.             return self.tree[tree_index]
  49.         mid = l + (r - l) // 2
  50.         left_tree_index = self.get_left_child(tree_index)
  51.         right_tree_index = self.get_right_child(tree_index)
  53.         if query_l >= (mid + 1):
  54.             return self.__query(right_tree_index, mid + 1, r, query_l, query_r)
  55.         elif query_r <= mid:
  56.             return self.__query(left_tree_index, l, mid, query_l, query_r)
  58.         left_res = self.__query(left_tree_index, l, mid, query_l, mid)
  59.         right_res = self.__query(right_tree_index, mid + 1, r, mid + 1, query_r)
  60.         return left_res + right_res
  62.     # 修改数据
  63.     def set(self, index, value):
  64.         self.data[index] = value
  65.         self.__set(0, 0, len(self.data)-1, index, value)
  67.     def __set(self, tree_index, l, r, index, value):
  68.         if l == r:
  69.             self.tree[tree_index] = value
  70.             return
  71.         mid = l + (r - l) // 2
  72.         left_tree_index = self.get_left_child(tree_index)
  73.         right_tree_index = self.get_right_child(tree_index)
  74.         if index >= mid + 1:
  75.             self.__set(right_tree_index, mid + 1, r, index, value)
  76.         else:
  77.             self.__set(left_tree_index, l, mid, index, value)
  78.         self.tree[tree_index] = self.tree[left_tree_index] + self.tree[right_tree_index]