How can I fix my code to solve for this puzzle? (Python)

Question

This is a diagram of the puzzle I'm solving for: Where A,B,C,D are four numbers that are either 0 or 1, and alphas are some random constants (could be positive or negative, but ignore the case when they're 0). The goal is to determine the list [A,B,C,D] given a random array of alpha. Here're some rules:

1. If alpha >0, then the neighboring points have different values. (Ex: if alpha_A = 4.5, then A,B = 01 or 10.

2. If alpha <0, then the neighboring points have the same value. (Ex: if alpha_B = -4.5, then B,C = 00 or 11.

3. Making the two adjacent points constant for the highest alpha_i, then start from the 2nd highest absolute value, iterate to the 2nd lowest one. Rule 1 and 2 could be violated only if alpha_i has the minimum absolute value.

Here's the current code I have:

alpha = np.random.normal(0, 10, N_object4)
pa = np.abs(alpha)
N_object4 = 4
num = pa.argsort()[-3:][::-1] # Index of descending order for abs
print(alpha)
print(num)
ABCD = np.zeros(N_object4).tolist()
if alpha[num[0]] > 0:
    ABCD[num[0]+1] = 1  # The largest assignment stays constant.
for i in range (1,len(num)): # Iterating from the 2nd largest abs(alpha) to the 2nd smallest.
    if i == 3:
        num[i+1] = num[0] 
    elif alpha[num[i]] > 0:
        if np.abs(num[i]-num[0]) == 1 or 2:
            ABCD[num[i+1]] = 1-ABCD[num[i]]
        elif np.abs(num[i]-num[0]) == 3:
            ABCD[num[i]] = 1-ABCD[num[0]]
    elif alpha[num[i]] < 0:
            if np.abs(num[i]-num[0]) == 1 or 2:
                **ABCD[num[i+1]] = ABCD[num[i]]**
            elif np.abs(num[i]-num[0]) == 3:
                ABCD[num[i]] = ABCD[num[0]]  

I think I'm still missing something in my code. There's an error in my current version (marked in the line with **):

IndexError: index 3 is out of bounds for axis 0 with size 3

I'm wondering where's my mistake?

Also, an example of the desired output: if

alpha = [12.74921599, -8.01870123, 11.07638142, -3.51723019] 

then

num = [0, 2, 1]

The correct list of ABCD should be:

ABCD = [0, 1, 1, 0] (or [1, 0, 0, 1])

My previous version could give me an output, but the list ABCD doesn't look right.

Thanks a lot for reading my question, and I really appreciate your help:)

Answer 1

Looking at your code if found a few smells:

• nums is 3 elements long, seems that you want num = pa.argsort()[-4:][::-1]

You also have indexing issues on other places:

if i == 3:
       num[i+1] = num[0] 

element 4 is not on the array

My personal advice: given the small size of the problem, get rid of numpy and just use python lists.

The other observation (that took me years to learn) is that for small problems, bruteforce is ok.

This is my bruteforce solution, with more time, this can be improved to do backtracking or dynamic programming:

import random
# To make things simpler, I will consider A = 0, and so on
pairs = [
    (0, 1),
    (1, 2),
    (2, 3),
    (3, 0),
]

# Ensure no zeros on our list
alphas = [(random.choice([1, -1]) * random.uniform(1, 10), pairs[n]) for n in range(4)]

print(alphas)

# Sort descending in
alphas.sort(reverse=True, key=lambda n: abs(n[0]))

print(alphas[:-1])


def potential_solutions():
    """Generator of potential solutions"""
    for n in range(16):
        # This just abuses that in binary, numbers from 0 to 15 cover all posible configurations
        # I use shift and masking to get each bit and return a list
        yield [n >> 3 & 1 , n >> 2 & 1, n >> 1 & 1, n & 1]


def check(solution):
    """Takes a candidate solution and check if is ok"""
    for alpha, (right, left) in alphas[:-1]:
        if alpha < 0 and solution[right] != solution[left]:
            return False
        elif alpha > 0 and solution[right] ==  solution[left]:
            return False
    return True


# Look for correct solutions
for solution in potential_solutions():
    if check(solution):
        print(solution)