Fitter Nearest-Neighbor function with three nested loops

Question

Background Problem

Consider a set of students that we want to place into Ng groups each of which contains Ns students. Let's label the students 1, 2, ... , Ng * Ns, and let an assignment of students be an array of shape (Ng, Ns) that encodes the groupings of students. For example, if Ng = 2 and Ns = 3, then an assignment of students could be:

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

note that any assignment that differs either in permuting students within groups, or in permuting the groups themselves would be considered an equivalent assignment. So for the assignment above, one equivalent assignment would be:

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

For a given assignment A, we define a neighboring assignment as any inequivalent assignment that differs from A by a single swap of students. For example, a neighbor of the first assignment above is obtained by swapping students 1 and 6 to obtain

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

Additionally, we have a fitness function f that takes an assignment A as its input, and outputs a real number f(A). The goal is to construct a function which I call fitter_neighbor which takes an assignment A and a fitness function f as its inputs and returns a fitter neighboring assignment provided such an assignment exists. If no such assignment exists, then fitter_neighbor should return A. It is not required that fitter_neighbor returns the fittest neighbor, and in fact, the goal is to have it systematically search through the neighbors and return the first fitter neighbor it comes across.

Current seemingly-not-very-pythonic code

def fitter_neighbor(A, f):

    Ng = len(A) #number of groups
    Ns = len(A[0]) #number of students per group
    Nn = int(Ns ** 2 * Ng * (Ng - 1) / 2) #number of neighboring assignments

    A_swap = np.copy(A) #initialize copy of assignment to be changed

    g1 = 0 #group number of person A in swap
    n_swaps = 0 #counter for number of swaps considered

    while g1 < Ng and f(A_swap) <= f(A) and n_swaps < Nn:
        s1 = 0
        while s1 < Ns and f(A_swap) <= f(A) and n_swaps < Nn:
            g2 = g1 + 1 #ensures that no swaps are repeated
            while g2 < Ng and f(A_swap) <= f(A) and n_swaps < Nn:
                s2 = 0
                while s2 < Ns and f(A_swap) <= f(A) and n_swaps < Nn:
                    A_swap = np.copy(A)
                    A_swap[g1, s1], A_swap[g2, s2] = A_swap[g2, s2], A_swap[g1, s1]
                    n_swaps += 1
                    s2 += 1
                g2 += 1    
            s1 += 1
        g1 += 1 
    if n_swaps < Nn:
        return A_swap
    else:
        return A

Can fitter_neighbor be written in a more pythonic way? In particular, is it possible to eliminate the nested while loops in some clever way?

As an example fitness function for the case Ns = 2, consider

def f2(A):
    return np.sum([(g[1] - g[0]) ** 2 for g in A])

Answer 1

Not a solution but a line of thought you may use for implementation. You can define the group of students as sets and then make tuples of these sets, so for example ({1, 3, 5}, {2, 4, 6}), let's call this tuple A and another tuple of sets B: (1, 4, 5}, {2, 3, 6}) I assume here that the order of the sets is important as you want to compare the first set of A with the first set of B, and so on...

You can then compare these tuples of sets by:

A = ({1, 3, 5}, {2, 4, 6}, {8, 10, 12})
B = ({1, 4, 5}, {2, 3, 6}, {8, 10, 12})
C = []
D = []

for i, a_list in enumerate(A):
   for j, b_list in enumerate(B):
       if i == j:
           C.append(a_list - b_list)
           D.append(b_list - a_list)

C = tuple(C)
D = tuple(D)
print(f'A: {A}\nB: {B}\nC: {C}\nD: {D}')

Answer:

A: ({1, 3, 5}, {2, 4, 6}, {8, 10, 12})
B: ({1, 4, 5}, {2, 3, 6}, {8, 10, 12})
C: ({3}, {4}, set())
D: ({4}, {3}, set())

C and D are then tuples with sets of differences for the first set, second set, etc. If these include sets with one element you know the sets in A and B are the same except for that one element.