|  | 
|  | 1 | +''' | 
|  | 2 | +Fenwick Tree implementation in Python. | 
|  | 3 | +
 | 
|  | 4 | +This module contains the implementation of Fenwick Trees, as well as two extensions of it: | 
|  | 5 | +- Range Update Point Query (RUPQ) | 
|  | 6 | +- Range Update Range Query (RURQ) | 
|  | 7 | +''' | 
|  | 8 | + | 
|  | 9 | +from typing import List | 
|  | 10 | + | 
|  | 11 | +class FenwickTree: | 
|  | 12 | +    ''' | 
|  | 13 | +    Fenwick Tree implementation in Python. | 
|  | 14 | +
 | 
|  | 15 | +    Attributes: | 
|  | 16 | +    - n: int - the size of the Fenwick Tree | 
|  | 17 | +    - ftree: list - the Fenwick Tree itself | 
|  | 18 | +
 | 
|  | 19 | +    Methods: | 
|  | 20 | +    - lsone(s: int) -> int: returns the least significant bit of s | 
|  | 21 | +    - query(i: int, j: int) -> int: returns the sum of the elements in the range [i, j] | 
|  | 22 | +    - update(i: int, v: int): updates the element at index i with value v | 
|  | 23 | +    - select(k: int) -> int: returns the index of the k-th element in the Fenwick Tree | 
|  | 24 | +    ''' | 
|  | 25 | +    def __init__(self, f: List[int]): | 
|  | 26 | +        self.n = len(f) | 
|  | 27 | +        self.ftree = [0] * (self.n + 1) | 
|  | 28 | + | 
|  | 29 | +        for i in range(1, self.n + 1): | 
|  | 30 | +            self.ftree[i] += f[i - 1] | 
|  | 31 | +            if i + self._lsone(i) <= self.n: | 
|  | 32 | +                self.ftree[i + self._lsone(i)] += self.ftree[i] | 
|  | 33 | + | 
|  | 34 | +    def _lsone(self, s: int) -> int: | 
|  | 35 | +        ''' | 
|  | 36 | +        Returns the least significant bit of s. | 
|  | 37 | +
 | 
|  | 38 | +        Args: | 
|  | 39 | +            s (int): The number to get the least significant bit from. | 
|  | 40 | +
 | 
|  | 41 | +        Returns: | 
|  | 42 | +            out (int): The least significant bit of s. | 
|  | 43 | +        ''' | 
|  | 44 | +        return s & -s | 
|  | 45 | + | 
|  | 46 | +    def query(self, i: int, j: int) -> int: | 
|  | 47 | +        ''' | 
|  | 48 | +        Queries the Fenwick Tree for the sum of the elements in the range [i, j]. | 
|  | 49 | +
 | 
|  | 50 | +        Args: | 
|  | 51 | +            i (int): The lower bound of the range. | 
|  | 52 | +            j (int): The upper bound of the range. | 
|  | 53 | +
 | 
|  | 54 | +        Returns: | 
|  | 55 | +            s (int): The sum of the elements in the range [i, j]. | 
|  | 56 | +        ''' | 
|  | 57 | +        if i > 1: | 
|  | 58 | +            return self.query(1, j) - self.query(1, i - 1) | 
|  | 59 | + | 
|  | 60 | +        s = 0 | 
|  | 61 | +        while j > 0: | 
|  | 62 | +            s += self.ftree[j] | 
|  | 63 | +            j -= self._lsone(j) | 
|  | 64 | + | 
|  | 65 | +        return s | 
|  | 66 | + | 
|  | 67 | +    def update(self, i: int, v: int): | 
|  | 68 | +        ''' | 
|  | 69 | +        Updates the element at index i with value v. | 
|  | 70 | +
 | 
|  | 71 | +        Args: | 
|  | 72 | +            i (int): The index of the element to update. | 
|  | 73 | +            v (int): The new value of the element. | 
|  | 74 | +        ''' | 
|  | 75 | + | 
|  | 76 | +        while i <= self.n: | 
|  | 77 | +            self.ftree[i] += v | 
|  | 78 | +            i += self._lsone(i) | 
|  | 79 | + | 
|  | 80 | +    def select(self, k: int) -> int: | 
|  | 81 | +        ''' | 
|  | 82 | +        Returns the index of the k-th element in the Fenwick Tree. | 
|  | 83 | +
 | 
|  | 84 | +        Args: | 
|  | 85 | +            k (int): The index of the element to return. | 
|  | 86 | +
 | 
|  | 87 | +        Returns: | 
|  | 88 | +            i (int): The index of the k-th element in the Fenwick Tree. | 
|  | 89 | +        ''' | 
|  | 90 | + | 
|  | 91 | +        p = 1 | 
|  | 92 | +        while p*2 <= self.n: | 
|  | 93 | +            p *= 2 | 
|  | 94 | + | 
|  | 95 | +        i = 0 | 
|  | 96 | +        while p > 0: | 
|  | 97 | +            if k > self.ftree[i + p]: | 
|  | 98 | +                k -= self.ftree[i + p] | 
|  | 99 | +                i += p | 
|  | 100 | +            p //= 2 | 
|  | 101 | + | 
|  | 102 | +        return i + 1 | 
|  | 103 | + | 
|  | 104 | +class RUPQ: | 
|  | 105 | +    ''' | 
|  | 106 | +    Range Update Point Query (RUPQ) implementation. | 
|  | 107 | +
 | 
|  | 108 | +    This class extends the Fenwick Tree to support range updates and point queries. | 
|  | 109 | +
 | 
|  | 110 | +    Attributes: | 
|  | 111 | +    - ftree: FenwickTree - the Fenwick Tree used for the range updates | 
|  | 112 | +
 | 
|  | 113 | +    Methods: | 
|  | 114 | +    - query(i: int) -> int: returns the value of the element at index i | 
|  | 115 | +    - update(i: int, j: int, v: int): updates the elements in the range [i, j] with value v | 
|  | 116 | +    ''' | 
|  | 117 | +    def __init__(self, n: int): | 
|  | 118 | +        self.ftree = FenwickTree([0] * n) | 
|  | 119 | + | 
|  | 120 | +    def query(self, i: int) -> int: | 
|  | 121 | +        ''' | 
|  | 122 | +        Queries the Fenwick Tree for the value of the element at index i. | 
|  | 123 | +         | 
|  | 124 | +        Args: | 
|  | 125 | +            i (int): The index of the element to query. | 
|  | 126 | +
 | 
|  | 127 | +        Returns: | 
|  | 128 | +            s (int): The value of the element at index i. | 
|  | 129 | +        ''' | 
|  | 130 | + | 
|  | 131 | +        return self.ftree.query(1, i) | 
|  | 132 | + | 
|  | 133 | +    def update(self, i: int, j: int, v: int): | 
|  | 134 | +        ''' | 
|  | 135 | +        Updates the elements in the range [i, j] with value v. | 
|  | 136 | +
 | 
|  | 137 | +        Args: | 
|  | 138 | +            i (int): The lower bound of the range. | 
|  | 139 | +            j (int): The upper bound of the range. | 
|  | 140 | +            v (int): The value to update the elements with. | 
|  | 141 | +        ''' | 
|  | 142 | + | 
|  | 143 | +        self.ftree.update(i, v) | 
|  | 144 | +        self.ftree.update(j + 1, -v) | 
|  | 145 | + | 
|  | 146 | +class RURQ: | 
|  | 147 | +    ''' | 
|  | 148 | +    Range Update Range Query (RURQ) implementation. | 
|  | 149 | +
 | 
|  | 150 | +    This class extends the Fenwick Tree to support range updates and range queries. | 
|  | 151 | +
 | 
|  | 152 | +    Attributes: | 
|  | 153 | +    - f: FenwickTree - the Fenwick Tree used for the range updates | 
|  | 154 | +    - r: RUPQ - the RUPQ used for the range queries | 
|  | 155 | +
 | 
|  | 156 | +    Methods: | 
|  | 157 | +    - query(i: int, j: int) -> int: returns the sum of the elements in the range [i, j] | 
|  | 158 | +    - update(i: int, j: int, v: int): updates the elements in the range [i, j] with value v | 
|  | 159 | +    ''' | 
|  | 160 | + | 
|  | 161 | +    def __init__(self, n: int): | 
|  | 162 | +        self.ftree = FenwickTree([0] * n) | 
|  | 163 | +        self.rupq = RUPQ(n) | 
|  | 164 | + | 
|  | 165 | +    def query(self, i: int, j: int) -> int: | 
|  | 166 | +        ''' | 
|  | 167 | +        Queries the Fenwick Tree for the sum of the elements in the range [i, j]. | 
|  | 168 | +
 | 
|  | 169 | +        Args: | 
|  | 170 | +            i (int): The lower bound of the range. | 
|  | 171 | +            j (int): The upper bound of the range. | 
|  | 172 | +
 | 
|  | 173 | +        Returns: | 
|  | 174 | +            s (int): The sum of the elements in the range [i, j]. | 
|  | 175 | +        ''' | 
|  | 176 | + | 
|  | 177 | +        if i > 1: | 
|  | 178 | +            return self.query(1, j) - self.query(1, i - 1) | 
|  | 179 | +        return self.rupq.query(j) * j - self.ftree.query(1, j) | 
|  | 180 | + | 
|  | 181 | +    def update(self, i: int, j: int, v: int): | 
|  | 182 | +        ''' | 
|  | 183 | +        Updates the elements in the range [i, j] with value v. | 
|  | 184 | +
 | 
|  | 185 | +        Args: | 
|  | 186 | +            i (int): The lower bound of the range. | 
|  | 187 | +            j (int): The upper bound of the range. | 
|  | 188 | +            v (int): The value to update the elements with. | 
|  | 189 | +        ''' | 
|  | 190 | + | 
|  | 191 | +        self.rupq.update(i, j, v) | 
|  | 192 | +        self.ftree.update(i, v * (i - 1)) | 
|  | 193 | +        self.ftree.update(j + 1, -1 * v * j) | 
0 commit comments