How to control the for loop layer numbers?

Viewed 83

I am a starter of Python and pretty interested in scripting with Python.

Recently, I faced a situation that I need to obtain all integer combination in the range of [-10, +10] and put them into two arrays in order, which means I have arrays [A] and [B], in [A], I need to have the first element to be [-10], the second [-10], and then in [B] it is [-10,-10], then I will erase the elements and make [A] have [-10,-10] but [B] have [-10, -9], then the next one is [A] = [-10,-10], [B] = [-10,-8]... until the last one [A] = [10,10], [B] = [10,10]. So I made 4-layer for loop to achieve that:

import numpy as np
A = np.zeros((1,2))
B = np.zeros((1,2))
for i in range(-10,+10+1):
    A[0][0] = I
    for ii in range(-10,+10+1):
        A[0][1] = ii
        for iii in range(-10,+10+1):
            B[0][0] = iii
            for iV in range(-10,+10+1):
                B[0][1] = iV

It works well to generate [A] and [B] in order in a stupid way. However, if I would like to change the size of [A] and [B], like A = np.zeros((1,3)), what should I do instead of adding more for loop lines? Is there a function to control the for loop numbers with different layers?

Thank you so much. Hope I have made my question clear LOL.

Best

Enxiao

2 Answers

You might want to try the itertools library:

from itertools import combinations_with_replacement
for combination in combinations_with_replacement(range(-10,11), n): # n being 4 in your case
    A[0] = combination[:n//2]
    B[0] = combination[n//2:]

The combinations_with_replacement generates tuples of n numbers in all combinations possible with the given numbers (range from -10 to 10); this can be split in two halves ([:n//2] -> first half, [n//2:] -> second half) and assigned to your array.

If you want to do it yourself without using itertools you can implement it like this:

import numpy as np

minVal, maxVal = -10, 10
sizeA, sizeB = 2, 2
totSize = sizeA+sizeB
AB = minVal * np.ones((totSize,))
while True:
    stop = True
    for i in range(totSize-1,-1,-1):
        if AB[i]<maxVal:
            AB[i] += 1
            AB[i+1:] = minVal
            stop = False
            break
    if stop:
        break

    A = AB[0:sizeA].reshape(1,-1)
    B = AB[sizeA:].reshape(1,-1)

    # Use A, B here
    print(A,B)
Related