How can I convert these repetitive if statements to a single loop in python?-CodePudding

I have this code:

#!/usr/bin/python3

def contract(e, i, c, n):
        l = len(e)
        grid = [[0 for i in range(i   1)] for x in range(l)]

        for num1, row1 in enumerate(grid):
                row1[0] = e[num1] #add exponents

        for num2, row2 in enumerate(grid):
                if 0 <= num2 < n[0]:
                        grid[num2][1] = c[num2]
                if n[0] <= num2 < n[0]   n[1]:
                        grid[num2][2] = c[num2]
                if n[0]   n[1] <= num2 < n[0]   n[1]   n[2]:
                        grid[num2][3] = c[num2]

        for g in grid:
                print(g)

e = [0, 1, 2, 3]
i = 3
c = [4, 5, 6, 7]
n = [1, 2, 1]

contract(e, i, c, n)

The idea of this code is that I have 2 dimensional grid that has dimensions len(e) x (i 1). The first column contains exponents e. The rest of the columns should contain coefficients c in such a way that n determines the positions of the coefficients in the grid. For example, since n[0] = 1, column 1, row 0 in the grid contains number 4. The next element in n is 2, so the next column in grid (column 2) should contain 2 numbers, meaning numbers 5 and 6 in rows below the row that I used previously (meaning rows 1 and 2 beause row 0 is already used). n[2] = 1 so grid[3][3] = 7, etc.

I implemented this with repetitive if-statements and the code works fine, the output is as it should be:

[0, 4, 0, 0]
[1, 0, 5, 0]
[2, 0, 6, 0]
[3, 0, 0, 7]

However, I would like to make a general code that can do this for any number of coefficients and exponents. How can I convert those repetitive if statements to a single loop?

CodePudding user response：

I would convert it into a for loop that keeps track of the sum of the elements seen so far, adjusting the corresponding element if the inequality holds for that iteration:

for num2, row2 in enumerate(grid):
    total = 0
    for n_idx, n_elem in enumerate(n):
        if total <= num2 < total   n_elem:
            grid[num2][n_idx   1] = c[num2]
        total  = n_elem

I would advise against using sum() in this loop, as it recomputes the sum from scratch on each iteration, which isn't very efficient.

CodePudding user response：

Use a loop that sums successive slices of the n list.

for num2, row2 in enumerate(grid):
    for idx in range(len(n)):
        if sum(n[:idx], 0) < num2 < sum(n[:idx 1], 0):
            grid[num2][idx 1] = c[num2]

This is a direct mapping of the code you wrote to a loop, and reasonable if n doesn't get too large. BrokenBenchmark's answer is optimized to take advantage of the fact that the sum of each slice is the sum of the previous slice plus the current element.