Home > Mobile >  Placing all row pairs in 2d numpy array into a 3d array
Placing all row pairs in 2d numpy array into a 3d array

Time:01-08

Consider the numpy 2d array shown:

a = [[8, 16, 4, 1, 0, 5],
     [3, 0, 0, 11, 9, 7],
     [5, 5, 8, 5, 15, 5],
     [2, 0, 2, 14, 2, 0],
     [0, 1, 2, 3, 4, 15]]

I want to find all row pairs from the array, placing them into a 3d array. When selecting row pairs, repeat rows are allowed, and the order of the 2 rows is not important. In the example, there are 15 such 2-row permutations, so I'm hoping to get a 3d array that is 15 layers deep:

     [[8, 16, 4, 1, 0, 5],
      [8, 16, 4, 1, 0, 5]],
     [[8, 16, 4, 1, 0, 5],
      [3, 0, 0, 11, 9, 7]],
     [[8, 16, 4, 1, 0, 5],
      [5, 5, 8, 5, 15, 5]],
...etc
     [[2, 0, 2, 14, 2, 0],
      [2, 0, 2, 14, 2, 0]],
     [[2, 0, 2, 14, 2, 0],
      [0, 1, 2, 3, 4, 15]],
     [[0, 1, 2, 3, 4, 15],
      [0, 1, 2, 3, 4, 15]]]

The actual starting arrays can be quite large, so I'm hoping an efficient solution could be suggested.

CodePudding user response:

I think you'll have to go outside of numpy to get all the pairs, for instance using itertools.combinations_with_replacement from the standard library. Example:

import itertools
import numpy as np

a = [[8, 16, 4, 1, 0, 5],
     [3, 0, 0, 11, 9, 7],
     [5, 5, 8, 5, 15, 5],
     [2, 0, 2, 14, 2, 0],
     [0, 1, 2, 3, 4, 15]]

out = np.array(list(itertools.combinations_with_replacement(a, 2)))

CodePudding user response:

You can use itertools.combinations_with_replacement to generate the pairs, then slice:

from itertools import combinations_with_replacement

N = 2
idx = np.array(list(combinations_with_replacement(range(a.shape[0]), r=N)))
out = a[idx]

Output:

array([[[ 8, 16,  4,  1,  0,  5],
        [ 8, 16,  4,  1,  0,  5]],

       [[ 8, 16,  4,  1,  0,  5],
        [ 3,  0,  0, 11,  9,  7]],

       [[ 8, 16,  4,  1,  0,  5],
        [ 5,  5,  8,  5, 15,  5]],

       [[ 8, 16,  4,  1,  0,  5],
        [ 2,  0,  2, 14,  2,  0]],

       [[ 8, 16,  4,  1,  0,  5],
        [ 0,  1,  2,  3,  4, 15]],

       [[ 3,  0,  0, 11,  9,  7],
        [ 3,  0,  0, 11,  9,  7]],

       [[ 3,  0,  0, 11,  9,  7],
        [ 5,  5,  8,  5, 15,  5]],

       [[ 3,  0,  0, 11,  9,  7],
        [ 2,  0,  2, 14,  2,  0]],

       [[ 3,  0,  0, 11,  9,  7],
        [ 0,  1,  2,  3,  4, 15]],

       [[ 5,  5,  8,  5, 15,  5],
        [ 5,  5,  8,  5, 15,  5]],

       [[ 5,  5,  8,  5, 15,  5],
        [ 2,  0,  2, 14,  2,  0]],

       [[ 5,  5,  8,  5, 15,  5],
        [ 0,  1,  2,  3,  4, 15]],

       [[ 2,  0,  2, 14,  2,  0],
        [ 2,  0,  2, 14,  2,  0]],

       [[ 2,  0,  2, 14,  2,  0],
        [ 0,  1,  2,  3,  4, 15]],

       [[ 0,  1,  2,  3,  4, 15],
        [ 0,  1,  2,  3,  4, 15]]])
  • Related