Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import pytest | |
| import ding | |
| ding.enable_numba = False # noqa | |
| from ding.utils import SumSegmentTree, MinSegmentTree # noqa | |
| class TestSumSegmentTree: | |
| def test_create(self): | |
| with pytest.raises(AssertionError): | |
| tree = SumSegmentTree(capacity=13) | |
| tree = SumSegmentTree(capacity=16) | |
| assert (tree.operation == 'sum') | |
| assert (tree.neutral_element == 0.) | |
| assert (max(tree.value) == 0.) | |
| assert (min(tree.value) == 0.) | |
| def test_set_get_item(self): | |
| tree = SumSegmentTree(capacity=4) | |
| elements = [1, 5, 4, 7] | |
| get_result = [] | |
| for idx, val in enumerate(elements): | |
| tree[idx] = val | |
| get_result.append(tree[idx]) | |
| assert (elements == get_result) | |
| assert (tree.reduce() == sum(elements)) | |
| assert (tree.reduce(0, 3) == sum(elements[:3])) | |
| assert (tree.reduce(0, 2) == sum(elements[:2])) | |
| assert (tree.reduce(0, 1) == sum(elements[:1])) | |
| assert (tree.reduce(1, 3) == sum(elements[1:3])) | |
| assert (tree.reduce(1, 2) == sum(elements[1:2])) | |
| assert (tree.reduce(2, 3) == sum(elements[2:3])) | |
| with pytest.raises(AssertionError): | |
| tree.reduce(2, 2) | |
| def test_find_prefixsum_idx(self): | |
| tree = SumSegmentTree(capacity=8) | |
| elements = [0, 0.1, 0.5, 0, 0, 0.2, 0.8, 0] | |
| for idx, val in enumerate(elements): | |
| tree[idx] = val | |
| with pytest.raises(AssertionError): | |
| tree.find_prefixsum_idx(tree.reduce() + 1e-4, trust_caller=False) | |
| with pytest.raises(AssertionError): | |
| tree.find_prefixsum_idx(-1e-6, trust_caller=False) | |
| assert (tree.find_prefixsum_idx(0) == 1) | |
| assert (tree.find_prefixsum_idx(0.09) == 1) | |
| assert (tree.find_prefixsum_idx(0.1) == 2) | |
| assert (tree.find_prefixsum_idx(0.59) == 2) | |
| assert (tree.find_prefixsum_idx(0.6) == 5) | |
| assert (tree.find_prefixsum_idx(0.799) == 5) | |
| assert (tree.find_prefixsum_idx(0.8) == 6) | |
| assert (tree.find_prefixsum_idx(tree.reduce()) == 6) | |
| class TestMinSegmentTree: | |
| def test_create(self): | |
| tree = MinSegmentTree(capacity=16) | |
| assert (tree.operation == 'min') | |
| assert (tree.neutral_element == np.inf) | |
| assert (max(tree.value) == np.inf) | |
| assert (min(tree.value) == np.inf) | |
| def test_set_get_item(self): | |
| tree = MinSegmentTree(capacity=4) | |
| elements = [1, -10, 10, 7] | |
| get_result = [] | |
| for idx, val in enumerate(elements): | |
| tree[idx] = val | |
| get_result.append(tree[idx]) | |
| assert (elements == get_result) | |
| assert (tree.reduce() == min(elements)) | |
| assert (tree.reduce(0, 3) == min(elements[:3])) | |
| assert (tree.reduce(0, 2) == min(elements[:2])) | |
| assert (tree.reduce(0, 1) == min(elements[:1])) | |
| assert (tree.reduce(1, 3) == min(elements[1:3])) | |
| assert (tree.reduce(1, 2) == min(elements[1:2])) | |
| assert (tree.reduce(2, 3) == min(elements[2:3])) | |