Skip to main content
2 of 2
Clarify title and question
Stephen Rauch
  • 4.2k
  • 12
  • 23
  • 36

Extended matrix generation

I am extending a matrix with some calculated values. This code runs very slowly.

Is there a way to vectorize this with slice and hstack?

import numpy as np

def expand(X):
    X_expanded = np.zeros((X.shape[0], 6))
    for i in range(X.shape[0]):
        X_expanded[i,0]=X[i,0]
        X_expanded[i,1]=X[i,1]
        X_expanded[i,2]=X[i,0]**2
        X_expanded[i,3]=X[i,1]**2
        X_expanded[i,4]=X[i,0]*X[i,1]
        X_expanded[i,5]=1
    return X_expanded