code: def test_prefixsum_idx():
      """
      test Segment Tree data structure
      """
      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
      assert np.all(
          tree.find_prefixsum_idx([0.0, 0.5, 0.99, 1.01, 3.00, 4.00])
          == [2, 2, 2, 3, 3, 3]
      )
  
      tree = SumSegmentTree(4)
  
      tree[np.array([2, 3])] = [1.0, 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
      assert np.all(
          tree.find_prefixsum_idx([0.0, 0.5, 0.99, 1.01, 3.00, 4.00])
          == [2, 2, 2, 3, 3, 3]
      )
  
imports: import numpy
  from stable_baselines.common.segment_tree import SumSegmentTree, MinSegmentTree
label_desc: Use image node_id 3 for calling a global function with example usage: test_prefixsum_idx() without return types