text: import numpy as np
  
  from baselines.common.segment_tree import SumSegmentTree, MinSegmentTree
  
  
  def test_tree_set():
      tree = SumSegmentTree(4)
  
      tree[2] = 1.0
      tree[3] = 3.0
  
      assert np.isclose(tree.sum(), 4.0)
      assert np.isclose(tree.sum(0, 2), 0.0)
      assert np.isclose(tree.sum(0, 3), 1.0)
      assert np.isclose(tree.sum(2, 3), 1.0)
      assert np.isclose(tree.sum(2, -1), 1.0)
      assert np.isclose(tree.sum(2, 4), 4.0)
  
  
  def test_tree_set_overlap():
      tree = SumSegmentTree(4)
  
      tree[2] = 1.0
      tree[2] = 3.0
  
      assert np.isclose(tree.sum(), 3.0)
      assert np.isclose(tree.sum(2, 3), 3.0)
      assert np.isclose(tree.sum(2, -1), 3.0)
      assert np.isclose(tree.sum(2, 4), 3.0)
      assert np.isclose(tree.sum(1, 2), 0.0)
  
  
  def test_prefixsum_idx():
      tree = SumSegmentTree(4)
  
      tree[2] = 1.0
      tree[3] = 3.0
  
      assert tree.find_prefixsum_idx(0.0) == 2
      assert tree.find_prefixsum_idx(0.5) == 2
      assert tree.find_prefixsum_idx(0.99) == 2
      assert tree.find_prefixsum_idx(1.01) == 3
      assert tree.find_prefixsum_idx(3.00) == 3
      assert tree.find_prefixsum_idx(4.00) == 3
  
  
  def test_prefixsum_idx2():
      tree = SumSegmentTree(4)
  
      tree[0] = 0.5
      tree[1] = 1.0
      tree[2] = 1.0
      tree[3] = 3.0
  
      assert tree.find_prefixsum_idx(0.00) == 0
      assert tree.find_prefixsum_idx(0.55) == 1
      assert tree.find_prefixsum_idx(0.99) == 1
      assert tree.find_prefixsum_idx(1.51) == 2
      assert tree.find_prefixsum_idx(3.00) == 3
      assert tree.find_prefixsum_idx(5.50) == 3
  
  
  def test_max_interval_tree():
      tree = MinSegmentTree(4)
  
      tree[0] = 1.0
      tree[2] = 0.5
      tree[3] = 3.0
  
      assert np.isclose(tree.min(), 0.5)
      assert np.isclose(tree.min(0, 2), 1.0)
      assert np.isclose(tree.min(0, 3), 0.5)
      assert np.isclose(tree.min(0, -1), 0.5)
      assert np.isclose(tree.min(2, 4), 0.5)
      assert np.isclose(tree.min(3, 4), 3.0)
  
      tree[2] = 0.7
  
      assert np.isclose(tree.min(), 0.7)
      assert np.isclose(tree.min(0, 2), 1.0)
      assert np.isclose(tree.min(0, 3), 0.7)
      assert np.isclose(tree.min(0, -1), 0.7)
      assert np.isclose(tree.min(2, 4), 0.7)
      assert np.isclose(tree.min(3, 4), 3.0)
  
      tree[2] = 4.0
  
      assert np.isclose(tree.min(), 1.0)
      assert np.isclose(tree.min(0, 2), 1.0)
      assert np.isclose(tree.min(0, 3), 1.0)
      assert np.isclose(tree.min(0, -1), 1.0)
      assert np.isclose(tree.min(2, 4), 3.0)
      assert np.isclose(tree.min(2, 3), 4.0)
      assert np.isclose(tree.min(2, -1), 4.0)
      assert np.isclose(tree.min(3, 4), 3.0)
  
  
  if __name__ == '__main__':
      test_tree_set()
      test_tree_set_overlap()
      test_prefixsum_idx()
      test_prefixsum_idx2()
      test_max_interval_tree()
  