Identifying elements with same "i" values in Python

June 20, 2022

I have an array A with dimensions (1,10,2). I want the code to give me the elements with the same i. If I have [0,1], i=0 and j=1. The desired output is attached.

import numpy as np
A = np.array([[[0, 1],
        [0, 2],
        [1, 3],
        [2, 3],
        [2, 5],
        [3, 4],
        [3, 6],
        [4, 7],
        [5, 6],
        [6, 7]]])

for i in range(0,len(A[0])):
    print([A[0,i]])

The desired output is

i=0: array([[[0, 1],[0,2]]])
i=1: array([[[1,3]]])
i=2: array([[[2, 3],[2,5]]])
i=3: array([[[3,4],[3,6]]])
i=4: array([[[4,7]]])
i=5: array([[[5,6]]])
i=6: array([[[6,7]]])

>Solution :

i_max = A[:,:,0].max()
for i in range(0,i_max+1):
    slice_as_str = str(A[A[:,:,0]==i]).replace("\n", "").replace(" ", ",")
    print(f"i={i}: array([{slice_as_str}])")