How to efficiently zip matrix elements with their row numbers in Python?-CodePudding

Suppose I have a 2D numpy array A:

2 7 6 
3 8 9

How can I do a kind of zip to get an array

(2,0) (7,0) (6,0)
(3,1) (8,1) (9,1)

So, elements in the matrix are zipped with the row number. What would be an efficient way to do it that performs well for a large matrix A?

CodePudding user response：

arr = np.array([[2, 7, 6], [3, 8, 9]])

row, col = arr.shape
res = np.dstack([arr, np.broadcast_to(np.arange(row), (col, row)).T])

assert np.array_equal(
    res,
    np.array(
        [
            [(2, 0), (7, 0), (6, 0)],
            [(3, 1), (8, 1), (9, 1)],
        ]
    ),
)

Performance Comparison

import numpy as np
import time

def f1(arr):
    row, col = arr.shape
    return np.dstack([arr, np.broadcast_to(np.arange(row), (col, row)).T])

def f2(arr):
    values = [(x, i[0]) for i, x in np.ndenumerate(arr)]
    return np.array(values, dtype="i4, i4").reshape(arr.shape)

def benchmark(f, *args):
    start = time.process_time()
    f(*args)
    return time.process_time() - start

if __name__ == "__main__":
    rg = np.random.default_rng(1)
    a = rg.random((3000, 3000))
    print(benchmark(f1, a))
    print(benchmark(f2, a))

0.046875
3.5625

CodePudding user response：

You can use numpy ndenumerate

arr = np.array([[2, 7, 6], [3, 8, 9]])
values = [(x, i[0]) for i, x in np.ndenumerate(arr)]
print(values)
# [(2, 0), (7, 0), (6, 0), (3, 1), (8, 1), (9, 1)]