[Leetcode] Binary Indexed Tree

Posted on | In , , |

Used to solve prefix sum problems. Why use BIT: can calculate prefix sums and update the sum with O(logn) efficiency.

Construct a BIT from a int array

LC collection: https://leetcode.com/tag/binary-indexed-tree/

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class BinaryIndexedTree:
    def __init__(self, nums):
        self.tree = [0] * (len(nums) + 1)
        # initialize tree
        for i in range(len(nums)):
            self.tree = self.update(i + 1, nums[i])

    def __lowBit(self, i):
        return i & -i

    def update(self, index, diff):
        # update tree[index] with value
        i = index
        while i < len(nums):
            self.tree[i] += diff
            i += self.lowBit(i) # get to upper level

    def prefixSum(self, index):
        res = 0
        i = index
        while i > 0:
            res += self.tree[i]
            i -= self.lowBit(i) # get next level sum
        return res

307. Range Sum Query - Mutable

Leetcode: https://leetcode.com/problems/range-sum-query-mutable/

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class NumArray:
    def __init__(self, nums: List[int]):
        self.nums = [0] * (len(nums))
        self.tree = [0] * (len(nums) + 1)
        # initialize
        for i in range(len(nums)):
            self.update(i, nums[i])

    def update(self, index: int, val: int) -> None:
        diff =  val - self.nums[index]
        self.nums[index] = val
        i = index + 1
        while i <= len(self.nums):
            self.tree[i] += diff
            i += i & -i
    
    def prefixSum(self, index):
        res = 0
        i = index + 1
        while i > 0:
            res += self.tree[i]
            i -= i & -i
        return res

    def sumRange(self, left: int, right: int) -> int:
        return self.prefixSum(right) - self.prefixSum(left - 1)

Binary Indexed Tree (Fenwick Array) (树状数组)

Indexed Tree: o(1) get child node and parent node

index i => complete tree
left child index = 2i + 1
right child index = 2i + 2
parent: (i - 1)/2

Use case:

  1. Range Sum, Query
  2. 所有BIT能解决的问题都能被 segment tree解决

implemented in array

315. Count of Smaller Numbers After Self

Leetcode: https://leetcode.com/problems/count-of-smaller-numbers-after-self/

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
# Got TLE somehow :(
def countSmaller(self, nums: List[int]) -> List[int]:
    size = 2 * 10 ** 4 + 2
    tree = [0] * size
    offset = 10 ** 4
        
    def update(i, diff):
        i += 1
        while i < len(tree):
            tree[i] += diff
            i += i & -i
            
    def prefixSum(i):
        res = 0
        while i > 0:
            res += tree[i]
            i -= i & -i
        return res
    res = []
    # initialize
    for i in reversed(range(len(nums))):
        index = nums[i] + offset
        res = [prefixSum(index)] + res
        update(index, 1)
        
    return res

327. Count of Range Sum

Leetcode: https://leetcode.com/problems/count-of-range-sum/