CODEWITHSUNDEEP
pythonlistloopsif-statement
Programmer
Programmer

Reputation: 324

Converting repetitive if statements into a loop

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 a 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 the 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 because 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 an extensible program that can do this for any number of coefficients and exponents. How can I convert those repetitive if statements to a single loop?

Upvotes: 1

Views: 106

Answers (2)

Barmar
Barmar

Reputation: 782564

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]) <= num2 < sum(n[:idx+1]):
            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.

Upvotes: 1

BrokenBenchmark
BrokenBenchmark

Reputation: 19262

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.

Upvotes: 1

Related Questions