mathdatech/Python
1
0
Fork 0
mirror of https://github.com/TheAlgorithms/Python.git synced 2025-02-26 02:48:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

[mypy] Fix type annotations in non_recursive_segment_tree (#5652)

Browse Source
This commit is contained in:
Dylan Buchi 2021-10-30 16:43:48 -03:00 committed by GitHub
parent e7565f8bfc
commit 678535b5c8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
data_structures/binary_tree/non_recursive_segment_tree.py
View File

@ -37,12 +37,12 @@ https://www.geeksforgeeks.org/segment-tree-efficient-implementation/
""" """
from __future__ import annotations from __future__ import annotations
from typing import Callable, TypeVar from typing import Any, Callable, Generic, TypeVar
T = TypeVar("T") T = TypeVar("T")
class SegmentTree: class SegmentTree(Generic[T]):
def __init__(self, arr: list[T], fnc: Callable[[T, T], T]) -> None: def __init__(self, arr: list[T], fnc: Callable[[T, T], T]) -> None:
""" """
Segment Tree constructor, it works just with commutative combiner. Segment Tree constructor, it works just with commutative combiner.
@ -55,8 +55,10 @@ class SegmentTree:
... lambda a, b: (a[0] + b[0], a[1] + b[1])).query(0, 2) ... lambda a, b: (a[0] + b[0], a[1] + b[1])).query(0, 2)
(6, 9) (6, 9)
""" """
self.N = len(arr) any_type: Any | T = None
self.st = [None for _ in range(len(arr))] + arr
self.N: int = len(arr)
self.st: list[T] = [any_type for _ in range(self.N)] + arr
self.fn = fnc self.fn = fnc
self.build() self.build()
@ -83,7 +85,7 @@ class SegmentTree:
p = p // 2 p = p // 2
self.st[p] = self.fn(self.st[p * 2], self.st[p * 2 + 1]) self.st[p] = self.fn(self.st[p * 2], self.st[p * 2 + 1])
def query(self, l: int, r: int) -> T: # noqa: E741 def query(self, l: int, r: int) -> T | None: # noqa: E741
""" """
Get range query value in log(N) time Get range query value in log(N) time
:param l: left element index :param l: left element index
@ -101,7 +103,8 @@ class SegmentTree:
7 7
""" """
l, r = l + self.N, r + self.N # noqa: E741 l, r = l + self.N, r + self.N # noqa: E741
res = None
res: T | None = None
while l <= r: # noqa: E741 while l <= r: # noqa: E741
if l % 2 == 1: if l % 2 == 1:
res = self.st[l] if res is None else self.fn(res, self.st[l]) res = self.st[l] if res is None else self.fn(res, self.st[l])
@ -135,7 +138,7 @@ if __name__ == "__main__":
max_segment_tree = SegmentTree(test_array, max) max_segment_tree = SegmentTree(test_array, max)
sum_segment_tree = SegmentTree(test_array, lambda a, b: a + b) sum_segment_tree = SegmentTree(test_array, lambda a, b: a + b)
def test_all_segments(): def test_all_segments() -> None:
""" """
Test all possible segments Test all possible segments
""" """