Scaling and incrementing non-zero elements of a NumPy matrix

I have a NumPy matrix C and want create a copy of it cPrime, which has some operation of the original matrix to all non-zero values of it. In the code below, for each non-zero element of C, I multiply by 30 and then add 1:

import numpy as np

size = 6
C = np.zeros((size,size), dtype=int)

C = 2
C = 5
C = 3
C = 1
C = 1
C = 1
C = 3

cPrime = np.zeros((size, size),dtype=int)
for i in range(size):
for j in range(size):
if C[i][j] != 0:
cPrime[i][j] = C[i][j]*30 + 1

This code works, but it feels inefficient. I'm about 99% sure that there's a really efficient way to achieve this goal, maybe using a masked array, but I haven't been able to figure it out.

migrated from programmers.stackexchange.comMar 21 '15 at 11:11

This question came from our site for professionals, academics, and students working within the systems development life cycle.

This is a typical use case for numpy.where:

cPrime = np.where(C, C * 30 + 1, C)

This is about twice as fast as (30 * C + 1) * (C != 0) and generalizes more easily to other conditions.

cPrime = 30 * C + 1, which uses broadcasting, comes close. Unfortunately, it adds 1 indiscriminately to every single element, even the elements that were originally zero.

C != 0 gives you the indexes of all the locations you want to operate on:

>>> C
array([[0, 0, 0, 1, 0, 0],
[0, 0, 0, 1, 0, 0],
[0, 0, 0, 1, 0, 0],
[0, 0, 0, 0, 0, 3],
[2, 5, 3, 0, 0, 0],
[0, 0, 0, 0, 0, 0]])
>>> C != 0
array([[False, False, False,  True, False, False],
[False, False, False,  True, False, False],
[False, False, False,  True, False, False],
[False, False, False, False, False,  True],
[ True,  True,  True, False, False, False],
[False, False, False, False, False, False]], dtype=bool)

Combining the two concepts, you can write either

cPrime = 30 * C + 1 * (C != 0)

or

cPrime = (30 * C + 1) * (C != 0)