How to Return The Steps After Merging 3 Sorted Arrays in Python

Question

I Just Want to Merge 3 Sorted Arrays in Python But Don't Want The Values Just Want The Steps Like From Where We Got The Values

Let's Say Here are the 3 Sorted Arrays

A = [1,3,4,6]
B = [2,3,4,5]
C = [1,5,9]

Original Merge Sort is

Sorted = [1,1,2,3,3,4,4,5,5,6,9]

But Instead of This I Want That From Which Array We Got The Value

Sorted = [A,C,B,A,B,A,B,B,C,A,C] 

You Could Say That I Want Steps That From Which Array We Got First & So On So Guide Me How We Can Do That on Python?

Answer 1

One approach:

from heapq import merge
from operator import itemgetter

A = [1, 3, 4, 6]
B = [2, 3, 4, 5]
C = [1, 5, 9]

lst = [[(val, label)for val in values] for label, values in zip(["A", "B", "C"], [A, B, C])]
res = [label for _, label in merge(*lst, key=itemgetter(0))]
print(res)

Output

['A', 'C', 'B', 'A', 'B', 'A', 'B', 'B', 'C', 'A', 'C']

First, for each input list create a list of tuples, where the first element of the tuple is the value and the second is a label indicating the origin, this is done in this line:

lst = [[(val, label) for val in values] for label, values in zip(["A", "B", "C"], [A, B, C])]

Afterwards use heapq.merge to merge the list, using as a key for the comparison only the first element of the tuple, then extract the label:

res = [label for _, label in merge(*lst, key=itemgetter(0))] 

The above code is equivalent to the following for-loop:

res = []
merged = merge(*lst, key=itemgetter(0))
for _, label in merged:
    res.append(label)