Iterate over i,j in a function returning two arrays-CodePudding

I am trying to accept two arrays from a function and iterate over value pairs in an array

import numpy as np

a = np.zeros(10).astype(np.uint8)
a[0:4] = 1
hist = np.zeros(4)

values, counts = np.unique(a, return_counts=True)
for u, c in zip(values, counts):
    hist[u]  = c
# This works. hist: [6. 4. 0. 0.]

# for u, c in zip(np.unique(a, return_counts=True)):  # ValueError: not enough values to unpack (expected 2, got 1)
#     hist[u]  = c
# for u, c in np.unique(a, return_counts=True):  # IndexError: index 6 is out of bounds for axis 0 with size 4
#     hist[u]  = c

Code works if I first accept two arrays, then use for k,v in zip(arr1, arr2) Is it possible two write for k,v in function_returning_two_arrays(args) as a one line statement?

Update. Both zip(*arg) and [arg] work. Can you please elaborate on this syntax, please. A link to an article would be enough. Then I can accept the answer. I got it that a * unpacks a tuple, but what does [some_tupple] do?

CodePudding user response：

Other than the unique step, this just basic python.

In [78]: a = np.zeros(10).astype(np.uint8)
    ...: a[0:4] = 1
    ...: ret = np.unique(a, return_counts=True)

unique returns a tuple of arrays, which can be used as is, or unpacked into 2 variables. I think unpacking makes the code clearer.

In [79]: ret
Out[79]: (array([0, 1], dtype=uint8), array([6, 4]))
In [80]: values, counts = ret
In [81]: values
Out[81]: array([0, 1], dtype=uint8)
In [82]: counts
Out[82]: array([6, 4])

The following just makes a list with 1 item - the tuple

In [83]: [ret]
Out[83]: [(array([0, 1], dtype=uint8), array([6, 4]))]

That's different from making a list of the two arrays - which just changes the tuple "wrapper" to a list:

In [84]: [values, counts]
Out[84]: [array([0, 1], dtype=uint8), array([6, 4])]

zip takes multiple items (it has a *args signature)

In [85]: list(zip(*ret))                 # same as zip(values, counts)
Out[85]: [(0, 6), (1, 4)]
In [86]: [(i,j) for i,j in zip(*ret)]    # using that in an iteration
Out[86]: [(0, 6), (1, 4)]
In [87]: [(i,j) for i,j in zip(values, counts)]
Out[87]: [(0, 6), (1, 4)]

So it pairs the nth element of values with the nth element of counts

Iteration on the [ret] list does something entirely different, or rather it does nothing - compare with `Out[83]:

In [88]: [(i,j) for i,j in [ret]]
Out[88]: [(array([0, 1], dtype=uint8), array([6, 4]))]

I think of list(zip(*arg)) as a list version of transpose:

In [90]: np.transpose(ret)
Out[90]: 
array([[0, 6],
       [1, 4]])
In [91]: [(i,j) for i,j in np.transpose(ret)]
Out[91]: [(0, 6), (1, 4)]