user9109129
user9109129

Reputation:

Avoiding multiple nested for-loops in python

How to avoid multiple nested for-loops when one nested for-loop has range up to the current iteration of the outer for-loop? For example, consider the following code: This program returns a triplet from a list arr such that arr[i] - arr[j] = arr[j] - arr[k] = d and i<j<k.

d =3
arr = [1, 2, 4, 5, 7, 8, 10]
list1 = []

for biggest in range(0, len(arr)):
    for bigger in range(0, biggest):
        for big in range(0, bigger):
            if abs(arr[big] - arr[bigger]) == d and abs(arr[bigger] - arr[biggest]) == d:
                list1.append([arr[big], arr[bigger], arr[biggest]])
print(list1))

Are there any other alternatives to using multiple nested loops?

Upvotes: 2

Views: 4929

Answers (4)

user9109129
user9109129

Reputation:

Found a better way I guess! avoiding the nested loops.

arr = [1,2,3,4,5,6,7,8,9]
d = 3
list1 = arr
list2 = []
for each in (0,len(list1)):
    if list1[each] + d in arr and list1[each] + 2*d in arr:
         list2.append([list1[each], list1[each]+d, list1[each]+2*d])
print(list2)

Upvotes: 1

Oluwafemi Sule
Oluwafemi Sule

Reputation: 38922

While previous answers are pythonic, if you care about the implementation of the search algorithm, you can reduce the complexity of your algorithm from O(N^3) to O(N^2logN) by implementing a binary search algorithm to find k in the space between j+1 and the length of lst which satisfies d - abs(lst[j] - lst[k]) == 0.

d = 3
lst = [1, 2, 4, 5, 7, 8, 10]


def bsearch(lst, left, right, j):
    while left < right:
        k = (left + right) // 2

        diff = d - abs(lst[j] - lst[k])

        if diff == 0:
            return k
        if diff > 0:
            left = k + 1
        else:
            right = k

    return None


l, result = len(lst), []
for i in range(l):
    for j in range(i + 1, l):
        diff = d - abs(lst[i] - lst[j])

        if diff != 0: continue
        k = bsearch(lst, j + 1, l, j)

        if not k: continue
        result.append((lst[i], lst[j], lst[k]))

print(result)

[(1, 4, 7), (2, 5, 8), (4, 7, 10)]

Upvotes: 0

Rory Daulton
Rory Daulton

Reputation: 22544

You can use the combinations function from itertools. Your code would then become:

from itertools import combinations

d = 3
arr = [1, 2, 4, 5, 7, 8, 10]
list1 = []
for big, bigger, biggest in combinations(arr, 3):
    if abs(big - bigger) == d and abs(bigger - biggest) == d:
        list1.append([big, bigger, biggest])

print(list1)

Which gives the same printout as your code (after you remove the extraneous right parenthesis on your last line):

[[1, 4, 7], [2, 5, 8], [4, 7, 10]]

Note that I changed the meanings of your variables big, bigger, biggest to be the array values rather than their indices. Working with values and avoiding indices is more pythonic and easier to understand.

You could also just do it in a list comprehension, for a slightly different look, avoiding the temporary list, and probable speed increase.

from itertools import combinations

d = 3
arr = [1, 2, 4, 5, 7, 8, 10]
print([[big, bigger, biggest]
        for big, bigger, biggest in combinations(arr, 3)
        if abs(big - bigger) == d and abs(bigger - biggest) == d
    ])

Upvotes: 3

Alex Hall
Alex Hall

Reputation: 36013

You can replace the three loops with:

from itertools import combinations

for big, bigger, biggest in combinations(range(0, len(arr)), 3):

You can replace all the code with:

print([t for t in combinations(arr, 3)
       if t[2] - t[1] == t[1] - t[0] == d])

Upvotes: 5

Related Questions