CODEWITHSUNDEEP
pythonlist-comprehension
Balines
Balines

Reputation: 23

How to up a counter if some condition is true in list comprehension?

I have this function: locs is a list of locations: [(0,0),(4,5)...] points is a list of values: [6,7,9...]

def foo(size,locs,points):
    matrix = []
    loc = 0
    rows = size[0]
    cols = size[1]
    for i in range(rows):
        row = []
        for j in range(cols):
            if ((i,j) in locs):
                row.append(points[loc])
                loc += 1
            else:  
                row.append(0)
        matrix.append(row)
    return matrix

It builds a matrix with 0s while inserting predefined points in the predefined location. I want to turn it into list comprehension. I have this:

loc = 0
rows = size[0]
cols = size[1]
matrix = [[points[loc] if (i,j) in locs else 0 for j in range(cols)] for i in range(rows)]

but I don't know how to update the loc to be loc + 1 if condition is true. I tried searching for this but couldn't find a result to try. Will gladly accept ideas. Thank you!

Upvotes: 2

Views: 90

Answers (1)

yatu
yatu

Reputation: 88236

If NumPy is an option, and since the task involves numerical matrix operations it does seems quite convenient, you have np.add.at to do exactly that:

a = np.zeros((rows, cols))
np.add.at(a, tuple(zip(*locs)), points)

Upvotes: 1

Related Questions