Nicolas H
Nicolas H

Reputation: 545

Python: adding index as new column to 2D array

Suppose I have np.array like below

dat = array([[ 0,  1,  0],
[ 1,  0,  0],
[0, 0, 1]]
)

What I want to do is that adding the (index of row + 1) as a new column to this array, which is like

newdat = array([[ 0,  1,  0, 1],
[ 1,  0,  0, 2],
[0, 0, 1, 3]]
)

How should I achieve this.

Upvotes: 4

Views: 1492

Answers (3)

Joe Ferndz
Joe Ferndz

Reputation: 8508

You can also use np.append(). You can also get more info about [...,None] here

import numpy as np

dat = np.array([
    [0, 1, 0],
    [1, 0, 0],
    [0, 0, 1]
])

a = np.array(range(1,4))[...,None] #None keeps (n, 1) shape
dat = np.append(dat, a, 1)

print (dat)

The output of this will be:

[[0 1 0 1]
 [1 0 0 2]
 [0 0 1 3]]

Or you can use hstack()

a = np.array(range(1,4))[...,None] #None keeps (n, 1) shape
dat = np.hstack((dat, a))

And as hpaulj mentioned, np.concatenate is the way to go. You can read more about concatenate documentation. Also, see additional examples of concatenate on stackoverflow

dat = np.concatenate([dat, a], 1)

Upvotes: 1

costaparas
costaparas

Reputation: 5237

Try something like this using numpy.insert():

import numpy as np

dat = np.array([
    [0, 1, 0],
    [1, 0, 0],
    [0, 0, 1]
])

dat = np.insert(dat, 3, values=[range(1, 4)], axis=1)

print(dat)

Output:

[[0 1 0 1]
 [1 0 0 2]
 [0 0 1 3]]

More generally, you can make use of numpy.ndarray.shape for the appropriate sizing:

dat = np.insert(dat, dat.shape[1], values=[range(1, dat.shape[0] + 1)], axis=1)

Upvotes: 0

Pablo C
Pablo C

Reputation: 4761

Use numpy.column_stack:

newdat = np.column_stack([dat, range(1,dat.shape[0] + 1)])
print(newdat)
#[[0 1 0 1]
# [1 0 0 2]
# [0 0 1 3]]

Upvotes: 1

Related Questions