I want to append elements of A (shape=(1,10,2)) with the same j to create a new array A1. For example, [1,3] and [2,3] should be appended into one element because of same j (=3) and different i (=1 and =2 respectively). The desired output is attached.

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

The desired output is

A1=array([[
        [0, 1],
        [0, 2],
        [[1, 3],[2, 3]],
        [2, 4],
        [3, 5],
        [[3, 6],[4, 6]],
        [[5, 7],[6, 7]]]])
A1.shape=(1,7,2)

CodePudding user response：

I've done it using the following steps. The only problem is that you can't have the final result as an array because of varying sizes. If you convert the result to a numpy array it becomes an array of lists of shape (7,).

You can however still iterate through it with for loops if it's not a huge list.

If you are using it in neural networks, you might want to consider converting to a ragged tensor

Get the list of second numbers

second_numbers = A[:,:,1].reshape(-1)

Get unique values from that list

uniq = set(second_numbers)

Create new list based on those unique values

new_list = []
for i in uniq:
  new_list.append((A[:, second_numbers == i, :].reshape(-1,2)).tolist())

Full code with result:

second_numbers = A[:,:,1].reshape(-1)
uniq = set(second_numbers)
new_list = []
for i in uniq:
  new_list.append((A[:, second_numbers == i, :].reshape(-1,2)).tolist())

new_list
>>> [[[0, 1]],
 [[0, 2]],
 [[1, 3], [2, 3]],
 [[2, 4]],
 [[3, 5]],
 [[3, 6], [4, 6]],
 [[5, 7], [6, 7]]]