Heaps and Priority Queues Questions

import numpy as np

class PrioritizedReplayBuffer:
    def __init__(self, capacity):
        self.capacity = int(capacity)
        # tree size: next power of two leaves simplifies indexing, but we'll use 2*capacity
        self.tree = np.zeros(2 * self.capacity, dtype=np.float64)  # 1-indexed conceptual: leaves at [capacity:capacity+N-1]
        self.data = [None] * self.capacity
        self.size = 0
        self.write = 0

    def _leaf_index(self, idx):
        return self.capacity + idx

    def add(self, experience, priority):
        # store data
        self.data[self.write] = experience
        # set priority
        leaf = self._leaf_index(self.write)
        self._update_tree(leaf, float(priority))
        # advance pointers
        self.write = (self.write + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)

    def _update_tree(self, leaf, value):
        # set leaf value then propagate differences up
        idx = int(leaf)
        old = self.tree[idx]
        self.tree[idx] = value
        diff = value - old
        idx //= 2
        while idx >= 1:
            self.tree[idx] += diff
            idx //= 2

    def update(self, indices, priorities):
        # indices: list/array of data indices (0..N-1)
        for i, p in zip(indices, priorities):
            leaf = self._leaf_index(int(i))
            self._update_tree(leaf, float(p))

    def total_priority(self):
        return float(self.tree[1])

    def sample(self, batch_size):
        assert self.size > 0
        batch = []
        idxs = []
        priorities = []
        total = self.total_priority()
        # stratified sampling for lower variance (optional)
        segment = total / batch_size
        for i in range(batch_size):
            a = segment * i
            b = segment * (i + 1)
            s = np.random.uniform(a, b)
            # traverse tree
            node = 1
            while node < self.capacity:
                left = 2 * node
                if self.tree[left] >= s:
                    node = left
                else:
                    s -= self.tree[left]
                    node = left + 1
            leaf_idx = node - self.capacity
            # clamp to current size (because capacity may be larger)
            if leaf_idx >= self.size:
                leaf_idx = self.size - 1
            batch.append(self.data[leaf_idx])
            idxs.append(leaf_idx)
            priorities.append(self.tree[node])
        return batch, idxs, np.array(priorities, dtype=np.float64)

import heapq
from typing import Iterator, List

def running_median(stream: Iterator[int]) -> Iterator[float]:
    """
    Yield running median after each integer from stream.
    left: max-heap implemented as negated min-heap (stores smaller half)
    right: min-heap (stores larger half)
    Invariant:
      - len(left) == len(right) or len(left) == len(right) + 1
      - all(x in left) <= all(y in right)
    """
    left: List[int] = []   # max-heap via negatives
    right: List[int] = []  # min-heap

    for x in stream:
        # push into appropriate heap
        if not left or x <= -left[0]:
            heapq.heappush(left, -x)
        else:
            heapq.heappush(right, x)

        # rebalance to maintain size invariant
        if len(left) > len(right) + 1:
            heapq.heappush(right, -heapq.heappop(left))
        elif len(right) > len(left):
            heapq.heappush(left, -heapq.heappop(right))

        # compute median
        if len(left) > len(right):
            yield float(-left[0])
        else:
            yield ( -left[0] + right[0] ) / 2.0