forked from kamyu104/LeetCode-Solutions
-
Notifications
You must be signed in to change notification settings - Fork 0
/
range-sum-query-mutable.py
133 lines (119 loc) · 4.51 KB
/
range-sum-query-mutable.py
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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
# Time: ctor: O(n),
# update: O(logn),
# query: O(logn)
# Space: O(n)
class NumArray(object):
def __init__(self, nums):
"""
initialize your data structure here.
:type nums: List[int]
"""
if not nums:
return
self.__nums = nums
self.__bit = [0] * (len(self.__nums) + 1)
for i in xrange(1, len(self.__bit)):
self.__bit[i] = nums[i-1] + self.__bit[i-1]
for i in reversed(xrange(1, len(self.__bit))):
last_i = i - (i & -i)
self.__bit[i] -= self.__bit[last_i]
def update(self, i, val):
"""
:type i: int
:type val: int
:rtype: int
"""
if val - self.__nums[i]:
self.__add(i, val - self.__nums[i])
self.__nums[i] = val
def sumRange(self, i, j):
"""
sum of elements nums[i..j], inclusive.
:type i: int
:type j: int
:rtype: int
"""
return self.__sum(j) - self.__sum(i-1)
def __sum(self, i):
i += 1
ret = 0
while i > 0:
ret += self.__bit[i]
i -= (i & -i)
return ret
def __add(self, i, val):
i += 1
while i <= len(self.__nums):
self.__bit[i] += val
i += (i & -i)
# Time: ctor: O(n),
# update: O(logn),
# query: O(logn)
# Space: O(n)
# Segment Tree solution.
class NumArray2(object):
def __init__(self, nums,
query_fn=lambda x, y: x+y,
update_fn=lambda x, y: y,
default_val=0):
"""
initialize your data structure here.
:type nums: List[int]
"""
N = len(nums)
self.__original_length = N
self.__tree_length = 2**(N.bit_length() + (N&(N-1) != 0))-1
self.__query_fn = query_fn
self.__update_fn = update_fn
self.__default_val = default_val
self.__tree = [default_val for _ in range(self.__tree_length)]
self.__lazy = [None for _ in range(self.__tree_length)]
self.__constructTree(nums, 0, self.__original_length-1, 0)
def update(self, i, val):
self.__updateTree(val, i, i, 0, self.__original_length-1, 0)
def sumRange(self, i, j):
return self.__queryRange(i, j, 0, self.__original_length-1, 0)
def __constructTree(self, nums, left, right, idx):
if left > right:
return
if left == right:
self.__tree[idx] = self.__update_fn(self.__tree[idx], nums[left])
return
mid = left + (right-left)//2
self.__constructTree(nums, left, mid, idx*2 + 1)
self.__constructTree(nums, mid+1, right, idx*2 + 2)
self.__tree[idx] = self.__query_fn(self.__tree[idx*2 + 1], self.__tree[idx*2 + 2])
def __apply(self, left, right, idx, val):
self.__tree[idx] = self.__update_fn(self.__tree[idx], val)
if left != right:
self.__lazy[idx*2 + 1] = self.__update_fn(self.__lazy[idx*2 + 1], val)
self.__lazy[idx*2 + 2] = self.__update_fn(self.__lazy[idx*2 + 2], val)
def __updateTree(self, val, range_left, range_right, left, right, idx):
if left > right:
return
if self.__lazy[idx] is not None:
self.__apply(left, right, idx, self.__lazy[idx])
self.__lazy[idx] = None
if range_left > right or range_right < left:
return
if range_left <= left and right <= range_right:
self.__apply(left, right, idx, val)
return
mid = left + (right-left)//2
self.__updateTree(val, range_left, range_right, left, mid, idx*2 + 1)
self.__updateTree(val, range_left, range_right, mid+1, right, idx*2 + 2)
self.__tree[idx] = self.__query_fn(self.__tree[idx*2 + 1],
self.__tree[idx*2 + 2])
def __queryRange(self, range_left, range_right, left, right, idx):
if left > right:
return self.__default_val
if self.__lazy[idx] is not None:
self.__apply(left, right, idx, self.__lazy[idx])
self.__lazy[idx] = None
if right < range_left or left > range_right:
return self.__default_val
if range_left <= left and right <= range_right:
return self.__tree[idx]
mid = left + (right-left)//2
return self.__query_fn(self.__queryRange(range_left, range_right, left, mid, idx*2 + 1),
self.__queryRange(range_left, range_right, mid + 1, right, idx*2 + 2))