Return indexes of the next sample in from prioritized distribution

def _sample_batch_prioritized(self, segment_tree, batch_size, history, forward_steps=1):
        """ Return indexes of the next sample in from prioritized distribution """
        p_total = segment_tree.total()
        segment = p_total / batch_size

        # Get batch of valid samples
        batch = [
            self._get_sample_from_segment(segment_tree, segment, i, history, forward_steps)
            for i in range(batch_size)
        ]

        probs, idxs, tree_idxs = zip(*batch)
        return np.array(probs), np.array(idxs), np.array(tree_idxs)