-
Notifications
You must be signed in to change notification settings - Fork 2
/
SumTree.py
50 lines (40 loc) · 1.41 KB
/
SumTree.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
import numpy as np
class SumTree:
def __init__(self, capacity):
self.write = 0
self.capacity = capacity
self.tree = np.zeros(2 * self.capacity - 1)
self.data = np.zeros(self.capacity, dtype=object)
def _propagate(self, idx, change):
parent = (idx - 1) // 2
self.tree[parent] += change
if parent != 0:
#print('parent', parent)
self._propagate(parent, change)
def _retreive(self, idx, s):
# since we count from 0
left = 2 * idx + 1
right = 2 * idx + 2
if left >= len(self.tree):
return idx
if s <= self.tree[left]:
return self._retreive(left, s)
else:
return self._retreive(right, s - self.tree[left])
def total(self):
return self.tree[0]
def add(self, priority, data):
idx = self.write + self.capacity - 1
self.data[self.write] = data
self.update(idx, priority)
self.write += 1
self.write %= self.capacity
def update(self, idx, priority):
change = priority - self.tree[idx]
#print('change is', change)
self.tree[idx] = priority
self._propagate(idx, change)
def get(self, s):
idx = self._retreive(0, s)
dataIdx = idx - self.capacity + 1
return self.data[dataIdx], idx