How to index a Python list with a Z3 Symbolic variable?

Question

I am trying to understand Python indexing on Symbolic variables using a SymbolicWalk. Here's a small snippet that I am unable to continue with:

from z3 import *

s = Solver()

Buckets = [3, 2, 5, 1, 4]

B_Weights = [7.3, 5.7, 4.25, 3.77, 6.02] # Weight of each bucket

B = len(Buckets)

pick = 3

X =[Int('B%d' % b) for b in range (pick)]

def pickBuckets(p):
    return Or([p == r for r in Buckets])

for p in range(pick):
    s.add(pickBuckets(X[p]))
    
# ------Calculate the weights of the picked buckets----
def SymbolicWalk(i, lst_B, lst_W, k):
#     Walk through B_Weights[i](lst_W[i]) and return the weight of the ith bucket

def SymbolicChoice(lst_B, lst_W):
#     select the index of X[p] in Buckets(lst_B) as i

def index(i, lst_B, lst_W):
    return SymbolicWalk(i, lst_B, lst_W, SymbolicChoice)

total_weight = 0
for p in range(pick):
    total_weight = total_weight + index(X[p], Buckets, B_Weights)
    
s.add(total_weight>15)
    

while s.check() == sat:
    m = s.model()
    picked = []
    for i in X:
        picked += [m[i]]

    print(picked)
    s.add(Or([p != v for p, v in zip(X, picked)]))

I should pick the index of each bucket selected in X[p] and then extract its weight from the list B_Weights. For example, if X[p]==5, then it should extract 4.25. One solution to the above example is:

[3, 1, 4]

This problem is probably a very simple form of symbolic walk but I'm still trying to figure out who it works to be able to solve a little more complicated problems.

Answer 1

This is a much simpler problem that you can solve simply by iterating over the pairs of values and adding the weights as appropriate:

from z3 import *

s = Solver()

Buckets   = [  3,   2,    5,    1,    4]
B_Weights = [7.3, 5.7, 4.25, 3.77, 6.02]

pick = 3
X = [Int('B%d' % b) for b in range (pick)]

s.add(Distinct(X))

total_weight = 0
for p in range(pick):
    s.add(Or([X[p] == b for b in Buckets]))
    total_weight = total_weight + sum([If(b == X[p], w, 0) for (b, w) in zip (Buckets, B_Weights)])

s.add(total_weight>15)

for s in all_smt(s, X):
   print(s)

where all_smt comes from https://stackoverflow.com/a/70656700/936310

When run, this produces:

[B0 = 2, B2 = 4, B1 = 1]
[B0 = 3, B2 = 1, B1 = 4]
[B0 = 1, B2 = 2, B1 = 4]
[B0 = 5, B2 = 2, B1 = 4]
[B0 = 4, B2 = 2, B1 = 1]
[B0 = 4, B2 = 5, B1 = 2]
[B0 = 4, B2 = 2, B1 = 5]
[B0 = 4, B2 = 2, B1 = 3]
...

plus a bunch of other solutions that satisfy the constraints.

Note that I added s.add(Distinct(X)) to make sure you don't pick the same bucket twice. In general, you probably want to ensure all the elements in X is ordered as well. (i.e., picking 4-2-3 is the same as 2-3-4.) If that's the case, you can add an ordered constraint on X's in a separate loop to reduce the solutions.