Python heapq.merge() Examples

def pauli_product(*elements: Pauli) -> Pauli:
    """Return the product of elements of the Pauli algebra"""
    result_terms = []

    for terms in product(*elements):
        coeff = reduce(mul, [term[1] for term in terms])
        ops = (term[0] for term in terms)
        out = []
        key = itemgetter(0)
        for qubit, qops in groupby(heapq.merge(*ops, key=key), key=key):
            res = next(qops)[1]  # Operator: X Y Z
            for op in qops:
                pair = res + op[1]
                res, rescoeff = PAULI_PROD[pair]
                coeff *= rescoeff
            if res != 'I':
                out.append((qubit, res))

        p = Pauli(((tuple(out), coeff),))
        result_terms.append(p)

    return pauli_sum(*result_terms) 

def tree_count_topologies(tree, sample_sets):
    for u in tree.samples():
        if not tree.is_leaf(u):
            raise ValueError("Internal samples not supported.")

    topology_counter = np.full(tree.tree_sequence.num_nodes, None, dtype=object)
    for sample_set_index, sample_set in enumerate(sample_sets):
        for u in sample_set:
            if not tree.is_sample(u):
                raise ValueError(f"Node {u} in sample_sets is not a sample.")
            topology_counter[u] = TopologyCounter.from_sample(sample_set_index)

    for u in tree.nodes(order="postorder"):
        children = []
        for v in tree.children(u):
            if topology_counter[v] is not None:
                children.append(topology_counter[v])
        if len(children) > 0:
            topology_counter[u] = combine_child_topologies(children)

    counters = []
    for root in tree.roots:
        if topology_counter[root] is not None:
            counters.append(topology_counter[root])
    return TopologyCounter.merge(counters) 

def _merged_breakpoints(
    fn0: PiecewiseConstantFunction[T],
    fn1: PiecewiseConstantFunction[T],
) -> Iterable[Tuple[XValue[T], float, float]]:

    bp0 = zip_longest(fn0.breakpoints.items(), [], fillvalue=0)
    bp1 = zip_longest(fn1.breakpoints.items(), [], fillvalue=1)
    yprev0, yprev1 = fn0._initial_value, fn1._initial_value

    for (x, y), fnnum in merge(bp0, bp1):
        if fnnum == 0:
            yield x, y, yprev1
            yprev0 = y
        elif fnnum == 1:
            yield x, yprev0, y
            yprev1 = y 

def execute(self):

        counter = 0
        for vcf_file_name in self.vcf_file_names:
            input_stream = InputStream(vcf_file_name, self.tempdir)

            samples = l_bp.parse_vcf(input_stream, self.vcf_lines, self.vcf_headers, include_ref=self.include_ref)
            for sample in samples:
                self.vcf_headers.append("##SAMPLE=<ID=" + sample + ">\n")
                self.has_genotypes = True
            counter += 1
            if counter > self.batchsize:
                self.vcf_lines.sort(key=l_bp.vcf_line_key)
                self.write_temp_file()
                counter = 0
        # no need to write the final batch to file
        # FIXME Replace this with a new VCF class with the headers all added
        self.write_header()

        self.vcf_lines.sort(key=l_bp.vcf_line_key)
        iterables = self.temp_files + [self.vcf_lines]
        self.output_handle.writelines(merge(*iterables))
        self.close_tempfiles() 

def _get_others_current_hand(self, player):
        player_up = self.players[get_upstream_player_id(player, self.players)]
        player_down = self.players[get_downstream_player_id(
            player, self.players)]
        others_hand = merge(player_up.current_hand, player_down.current_hand, key=functools.cmp_to_key(doudizhu_sort_card))
        return cards2str(others_hand)

#if __name__ == '__main__':
#    import numpy as np
#    game = SimpleDoudizhuGame()
#    state, player_id = game.init_game()
#    print(state)
#    while not game.is_over():
#        action = np.random.choice(list(state['actions']))
#        print(action)
#        state, next_player_id = game.step(action)
#        print(state) 

def get_prescriptions_for_dates(dates):
    """
    Yield all prescribing data for the given dates as tuples of the form:

        bnf_code, practice_code, date, items, quantity, actual_cost, net_cost

    sorted by bnf_code, practice and date.
    """
    dates = sorted(dates)
    filenames = [get_prescribing_filename(date) for date in dates]
    missing_files = [f for f in filenames if not os.path.exists(f)]
    if missing_files:
        raise RuntimeError(
            "Some required CSV files were missing:\n  {}".format(
                "\n  ".join(missing_files)
            )
        )
    prescribing_streams = [read_gzipped_prescribing_csv(f) for f in filenames]
    # We assume that the input files are already sorted by (bnf_code, practice,
    # month) so to ensure that the combined stream is sorted we just need to
    # merge them correctly, which heapq.merge handles nicely for us
    return heapq.merge(*prescribing_streams) 

def nthSuperUglyNumber(self, n, primes):
    uglies = [1]

    def gen(prime):
        for ugly in uglies:
            yield ugly * prime

    merged = heapq.merge(*map(gen, primes))
    while len(uglies) < n:
        ugly = next(merged)
        if ugly != uglies[-1]:
            uglies.append(ugly)
    return uglies[-1]


# [23] https://leetcode.com/problems/merge-k-sorted-lists/
# Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.
#
# use merge 

def _placement_points_generator(self, skyline, width):
        """Returns a generator for the x coordinates of all the placement
        points on the skyline for a given rectangle.

        WARNING: In some cases could be duplicated points, but it is faster
        to compute them twice than to remove them.
        
        Arguments:
            skyline (list): Skyline HSegment list
            width (int, float): Rectangle width

        Returns:
            generator
        """ 
        skyline_r = skyline[-1].right
        skyline_l = skyline[0].left

        # Placements using skyline segment left point
        ppointsl = (s.left for s in skyline if s.left+width <= skyline_r)

        # Placements using skyline segment right point
        ppointsr = (s.right-width for s in skyline if s.right-width >= skyline_l)

        # Merge positions
        return heapq.merge(ppointsl, ppointsr) 

def __init__(self, width, height, rot=True, *args, **kwargs):
        """
        _skyline is the list used to store all the skyline segments, each 
        one is a list with the format [x, y, width] where x is the x
        coordinate of the left most point of the segment, y the y coordinate
        of the segment, and width the length of the segment. The initial 
        segment is allways [0, 0, surface_width]
        
        Arguments:
            width (int, float): 
            height (int, float):
            rot (bool): Enable or disable rectangle rotation
        """
        self._waste_management = False
        self._waste = WasteManager(rot=rot)
        super(Skyline, self).__init__(width, height, rot, merge=False, *args, **kwargs) 

def _merge_catchup_txns(existing_txns, new_txns):
        """
        Merge any newly received txns during catchup with already received txns
        :param existing_txns:
        :param new_txns:
        :return:
        """
        # TODO: Can we replace this with SortedDict and before merging substract existing transactions from new?
        idx_to_remove = []
        start_seq_no = new_txns[0][0]
        end_seq_no = new_txns[-1][0]
        for seq_no, _ in existing_txns:
            if seq_no < start_seq_no:
                continue
            if seq_no > end_seq_no:
                break
            idx_to_remove.append(seq_no - start_seq_no)
        for idx in reversed(idx_to_remove):
            new_txns.pop(idx)

        return list(merge(existing_txns, new_txns, key=lambda v: v[0])) 

def mergeKLists(lists):
    # create a generator
    def gen(node):
        while node:
            yield node.val, node
            node = node.next

    dummy = last = ListNode(None)
    for _, last.next in heapq.merge(*map(gen, lists), key=lambda x: x[0]):
        last = last.next
    return dummy.next


# [23] https://leetcode.com/problems/merge-k-sorted-lists/
# Merge k sorted linked lists and return it as one sorted list. Analyze and describe its complexity.
#
# use push and pop 

def paulis_commute(element0: Pauli, element1: Pauli) -> bool:
    """
    Return true if the two elements of the Pauli algebra commute.
    i.e. if element0 * element1 == element1 * element0

    Derivation similar to arXiv:1405.5749v2 for the check_commutation step in
    the Raesi, Wiebe, Sanders algorithm (arXiv:1108.4318, 2011).
    """

    def _coincident_parity(term0: PauliTerm, term1: PauliTerm) -> bool:
        non_similar = 0
        key = itemgetter(0)

        op0 = term0[0]
        op1 = term1[0]
        for _, qops in groupby(heapq.merge(op0, op1, key=key), key=key):

            listqops = list(qops)
            if len(listqops) == 2 and listqops[0][1] != listqops[1][1]:
                non_similar += 1
        return non_similar % 2 == 0

    for term0, term1 in product(element0, element1):
        if not _coincident_parity(term0, term1):
            return False

    return True 

def _get_others_current_hand(self, player):
        player_up = self.players[get_upstream_player_id(player, self.players)]
        player_down = self.players[get_downstream_player_id(
            player, self.players)]
        others_hand = merge(player_up.current_hand, player_down.current_hand, key=functools.cmp_to_key(doudizhu_sort_card))
        return cards2str(others_hand)

#if __name__ == '__main__':

    # test init game
    #game = DoudizhuGame()
    #game.init_game() 

def batch_sort(input_iterator, output_path, buffer_size=32000, output_class=None):
    """batch sort helper with temporary files, supports sorting large stuff"""
    if not output_class:
        output_class = input_iterator.__class__

    chunks = []
    try:
        while True:
            current_chunk = list(islice(input_iterator,buffer_size))
            if not current_chunk:
                break
            current_chunk.sort()
            output_chunk = output_class(os.path.join(TMPD, "%06i" % len(chunks)))
            chunks.append(output_chunk)

            for elem in current_chunk:
                output_chunk.write(elem.obj)
            output_chunk.close()

        output_file = output_class(output_path)
        for elem in heapq.merge(*chunks):
            output_file.write(elem.obj)
        output_file.close()
    finally:
        for chunk in chunks:
            try:
                chunk.close()
                os.remove(chunk.name)
            except Exception:
                pass

# magic 

def merge(*iterables):
   keyed_iterables = [(Keyed(l_bp.vcf_line_key(obj), obj) for obj in iterable)
                        for iterable in iterables]
   for element in heapq.merge(*keyed_iterables):
       yield element.obj 

def posts_for_user(user: User, limit: Optional[int] = None) -> List[Post]:
    relevant = merge(*[user_posts[u] for u in following[user]], reverse=True)
    return list(islice(relevant, limit)) 

def finalize(self):
        if self.tmp_sdb is None:
            return
        self.finish_bucket()
        num_samples = len(self.tmp_sdb)
        self.tmp_sdb.close()
        self.tmp_sdb = None
        if self.buffered_samples is None:
            avg_sample_size = self.overall_size / max(1, num_samples)
            max_cached_samples = self.cache_size / max(1, avg_sample_size)
            buffer_size = max(1, int(max_cached_samples / max(1, len(self.buckets))))
        else:
            buffer_size = self.buffered_samples
        sdb_reader = SDB(self.tmp_sdb_filename, buffering=self.buffering, id_prefix='#pre-sorted')

        def buffered_view(bucket):
            start, end = bucket
            buffer = []
            current_offset = start
            while current_offset < end:
                while len(buffer) < buffer_size and current_offset < end:
                    buffer.insert(0, sdb_reader[current_offset])
                    current_offset += 1
                while len(buffer) > 0:
                    yield buffer.pop(-1)

        bucket_views = list(map(buffered_view, self.buckets))
        interleaved = heapq.merge(*bucket_views, key=lambda s: s.duration)
        with DirectSDBWriter(self.sdb_filename,
                             buffering=self.buffering,
                             audio_type=self.audio_type,
                             id_prefix=self.id_prefix) as sdb_writer:
            for index, sample in enumerate(interleaved):
                old_id = sample.sample_id
                sdb_writer.add(sample)
                self.meta_list.append(self.meta_dict[old_id])
                del self.meta_dict[old_id]
                yield index / num_samples
        sdb_reader.close()
        os.unlink(self.tmp_sdb_filename) 

def __iter__(self):
        return heapq.merge(*self.iterables, key=self.key) 

def utilization_queue(self, free_capacity, visitor=None):
        """Returns utilization queue including the sub-allocs.

        All app queues from self and sub-allocs are merged in standard order,
        and then utilization is recalculated based on total reserved capacity
        of this alloc and sub-allocs combined.

        The function maintains invariant that any app (self or inside sub-alloc
        with utilization < 1 will remain with utilzation < 1.
        """
        total_reserved = self.total_reserved()
        queues = [
            alloc.utilization_queue(free_capacity, visitor)
            for alloc in six.itervalues(self.sub_allocations)
        ]

        queues.append(self.priv_utilization_queue())

        acc_demand = zero_capacity()
        available = total_reserved + free_capacity + np.finfo(float).eps
        util_before = utilization(acc_demand, total_reserved, available)

        for item in heapq.merge(*queues):
            rank, _u_before, _u_after, pending, order, app = item
            acc_demand = acc_demand + app.demand
            util_after = utilization(acc_demand, total_reserved, available)
            if app.priority == 0:
                util_before = _MAX_UTILIZATION
                util_after = _MAX_UTILIZATION

            # - lower rank allocations take precedence.
            # - for same rank, utilization takes precedence
            # - False < True, so for apps with same utilization we prefer
            #   those that already running (False == not pending)
            # - Global order
            entry = (rank, util_before, util_after, pending, order, app)
            if visitor:
                visitor(self, entry, acc_demand)

            util_before = util_after
            yield entry 

def get_req_data(self, req_type):
        if req_type is None:
            successes = [successes for (successes, _) in self.data_by_type.values()]
            failures = [failures for (_, failures) in self.data_by_type.values()]
            return (
                list(heapq.merge(*successes)),
                list(heapq.merge(*failures)),
            )

        return self.data_by_type[req_type] 

def create_accumulator(self) -> Histogram:
    # The number of accumulator (histogram) buckets is variable and depends on
    # the number of distinct intervals that are matched during calls to
    # add_inputs. This allows the historam size to start small and gradually
    # grow size during calls to merge until reaching the final histogram.
    return [] 

def merge_accumulators(self, accumulators: List[Histogram]) -> Histogram:
    result = []
    for bucket_id, buckets in itertools.groupby(
        heapq.merge(*accumulators), key=operator.attrgetter('bucket_id')):
      total_weighted_labels = 0.0
      total_weighted_predictions = 0.0
      total_weighted_examples = 0.0
      for bucket in buckets:
        total_weighted_labels += bucket.weighted_labels
        total_weighted_predictions += bucket.weighted_predictions
        total_weighted_examples += bucket.weighted_examples
      result.append(
          Bucket(bucket_id, total_weighted_labels, total_weighted_predictions,
                 total_weighted_examples))
    return result 

def merge(*iterables):
    # based on code posted by Scott David Daniels in c.l.p.
    # http://groups.google.com/group/comp.lang.python/msg/484f01f1ea3c832d

    keyed_iterables = [
        (Keyed(keyfunc(obj), obj)
        for obj in iterable)
        for iterable in iterables]
    for element in heapq.merge(*keyed_iterables):
        yield element.obj 

def batch_sort(infile, output, buffer_size=32000, tempdirs=None):
    if tempdirs is None:
        tempdirs = []
    if not tempdirs:
        tempdirs.append(gettempdir())

    chunks = []
    try:
        with io.open(infile, mode='r', buffering=64*1024, encoding='utf8') as input_file:
            print(u"Opened input {0}".format(infile))
            input_iterator = iter(input_file)
            for tempdir in cycle(tempdirs):
                current_chunk = list(islice(input_iterator,buffer_size))
                if not current_chunk:
                    break
                current_chunk.sort(key=keyfunc)
                fname = '%06i' % len(chunks)
                output_chunk = io.open(os.path.join(tempdir,fname),mode='w+',buffering=64*1024, encoding='utf8')
                print(u"Writing tempfile {0}/{1}".format(tempdir, fname))
                chunks.append(output_chunk)
                output_chunk.writelines(current_chunk)
                output_chunk.flush()
                output_chunk.seek(0)
        print(u"Writing outfile {0}".format(output))
        with io.open(output,mode='w',buffering=64*1024, encoding='utf8') as output_file:
            output_file.writelines(merge(*chunks))
    finally:
        for chunk in chunks:
            # noinspection PyBroadException
            try:
                chunk.close()
                os.remove(chunk.name)
            except:
                print(u"Exception when closing chunk")
                pass 

def collate(*iterables, **kwargs):
    """Return a sorted merge of the items from each of several already-sorted
    *iterables*.

        >>> list(collate('ACDZ', 'AZ', 'JKL'))
        ['A', 'A', 'C', 'D', 'J', 'K', 'L', 'Z', 'Z']

    Works lazily, keeping only the next value from each iterable in memory. Use
    :func:`collate` to, for example, perform a n-way mergesort of items that
    don't fit in memory.

    If a *key* function is specified, the iterables will be sorted according
    to its result:

        >>> key = lambda s: int(s)  # Sort by numeric value, not by string
        >>> list(collate(['1', '10'], ['2', '11'], key=key))
        ['1', '2', '10', '11']


    If the *iterables* are sorted in descending order, set *reverse* to
    ``True``:

        >>> list(collate([5, 3, 1], [4, 2, 0], reverse=True))
        [5, 4, 3, 2, 1, 0]

    If the elements of the passed-in iterables are out of order, you might get
    unexpected results.

    On Python 2.7, this function delegates to :func:`heapq.merge` if neither
    of the keyword arguments are specified. On Python 3.5+, this function
    is an alias for :func:`heapq.merge`.

    """
    if not kwargs:
        return merge(*iterables)

    return _collate(*iterables, **kwargs)


# If using Python version 3.5 or greater, heapq.merge() will be faster than
# collate - use that instead. 

def sorted_items(iterables):
    sorted_iterable = heapq.merge(*(_decorate_items(s) for s in iterables))

    for _, item in sorted_iterable:
        yield item 

def data_gen(self):
        yield from heapq.merge(*self.readers) 

def apply(func, *mappings):
    """Return a new IntervalMapping applying func to all values of mappings.
    
    For example, apply(lambda x, y: x + y, m1, m2) returns a new
    matching with the sum of values for m1 and
    m2. IntervalMapping.nothing is passed to func when the value is
    undefined.

    >>> m1 = IntervalMapping()
    >>> m2 = IntervalMapping()
    >>> m1[:] = m2[:] = 0 # avoid problems with undefined values
    >>> m1[0:2] = 1
    >>> m2[1:3] = 2
    >>> m3 = apply(lambda a, b: a + b, m1, m2)
    >>> m3[-1], m3[0], m3[1], m3[2], m3[3]
    (0, 1, 3, 2, 0)
    """

    values = [m.leftmost() for m in mappings]

    def changes():

        def start_i_value(i_m):
            i, m = i_m
            res = ((k.start, i, v) for k, v in m.iteritems(True))
            next(res)
            return res
        change_points = merge(*map(start_i_value, enumerate(mappings)))

        lastbound = None
        for bound, i, v in change_points:
            values[i] = v
            if bound != lastbound:
                yield bound, func(*values)
                lastbound = bound
        yield bound, func(*values)

    return IntervalMapping.from_changes(func(*values), changes()) 

def collate(*iterables, **kwargs):
    """Return a sorted merge of the items from each of several already-sorted
    *iterables*.

        >>> list(collate('ACDZ', 'AZ', 'JKL'))
        ['A', 'A', 'C', 'D', 'J', 'K', 'L', 'Z', 'Z']

    Works lazily, keeping only the next value from each iterable in memory. Use
    :func:`collate` to, for example, perform a n-way mergesort of items that
    don't fit in memory.

    If a *key* function is specified, the iterables will be sorted according
    to its result:

        >>> key = lambda s: int(s)  # Sort by numeric value, not by string
        >>> list(collate(['1', '10'], ['2', '11'], key=key))
        ['1', '2', '10', '11']


    If the *iterables* are sorted in descending order, set *reverse* to
    ``True``:

        >>> list(collate([5, 3, 1], [4, 2, 0], reverse=True))
        [5, 4, 3, 2, 1, 0]

    If the elements of the passed-in iterables are out of order, you might get
    unexpected results.

    On Python 3.5+, this function is an alias for :func:`heapq.merge`.

    """
    if not kwargs:
        return merge(*iterables)

    return _collate(*iterables, **kwargs)


# If using Python version 3.5 or greater, heapq.merge() will be faster than
# collate - use that instead. 

def collate(*iterables, **kwargs):
    """Return a sorted merge of the items from each of several already-sorted
    *iterables*.

        >>> list(collate('ACDZ', 'AZ', 'JKL'))
        ['A', 'A', 'C', 'D', 'J', 'K', 'L', 'Z', 'Z']

    Works lazily, keeping only the next value from each iterable in memory. Use
    :func:`collate` to, for example, perform a n-way mergesort of items that
    don't fit in memory.

    If a *key* function is specified, the iterables will be sorted according
    to its result:

        >>> key = lambda s: int(s)  # Sort by numeric value, not by string
        >>> list(collate(['1', '10'], ['2', '11'], key=key))
        ['1', '2', '10', '11']


    If the *iterables* are sorted in descending order, set *reverse* to
    ``True``:

        >>> list(collate([5, 3, 1], [4, 2, 0], reverse=True))
        [5, 4, 3, 2, 1, 0]

    If the elements of the passed-in iterables are out of order, you might get
    unexpected results.

    On Python 2.7, this function delegates to :func:`heapq.merge` if neither
    of the keyword arguments are specified. On Python 3.5+, this function
    is an alias for :func:`heapq.merge`.

    """
    if not kwargs:
        return merge(*iterables)

    return _collate(*iterables, **kwargs)


# If using Python version 3.5 or greater, heapq.merge() will be faster than
# collate - use that instead.