The elementwise square of subset of numpy array

Question

I have the following numpy array

np.array([[1, 2, 3],
          [4, 5, 6],
          [7, 8, 9]])

I want to compute the elementwise square of second column and third column only by retaining the first column as it is, yielding the result,

np.array([[1,  4,   9],
          [4, 25, 36]],
          [7, 64, 81]])

What have I tried I extracted the first column. then I extracted the second and third columns. found square, using numpy.square function.

arr1 = arr[:, 0]
arr2 = np.square(arr[:, 1:])

and then concatenated them

np.c_[arr1, arr2]

Is there a single step solution?

Answer 1

Select all elements of the non-first column using [:, 1:].

arr[:, 1:] **= 2