How to construct a complex constraint in PULP python

Question

I'm trying to solve a napsack-style optimization problem with additional complexity.

Here is a simple example. I'm trying to select 5 items that maximize value. 2 of the items must be orange, one must be blue, one must be yellow, and one must be red. This is straightforward. However, I want to add a constraint that the selected yellow, red, and orange items can only have one shape in common with the selected blue item.

import pandas as pd
from pulp import *
import re

BlueMatch = lambda x: 1 if x=='blue' else 0
YellowMatch = lambda x: 1 if x=='yellow' else 0
RedMatch = lambda x: 1 if x=='red' else 0
OrangeMatch = lambda x: 1 if x=='orange' else 0

data = [['A', 'blue', 'circle', 0.454], 
    ['B', 'yellow', 'square', 0.570],
    ['C', 'red', 'triangle', 0.789],
    ['D', 'red', 'circle', 0.718],
    ['E', 'red', 'square', 0.828],
    ['F', 'orange', 'square', 0.709],
    ['G', 'blue', 'circle', 0.696],
    ['H', 'orange', 'square', 0.285],
    ['I', 'orange', 'square', 0.698],
    ['J', 'orange', 'triangle', 0.861],
    ['K', 'blue', 'triangle', 0.658],
    ['L', 'yellow', 'circle', 0.819],
    ['M', 'blue', 'square', 0.352],
    ['N', 'orange', 'circle', 0.883],
    ['O', 'yellow', 'triangle', 0.755]]

example = pd.DataFrame(data, columns = ['item', 'color', 'shape', 'value']) 

example['color'] = example['color'].astype(str)

example['isBlue'] = example.color.apply(BlueMatch)
example['isYellow'] = example.color.apply(YellowMatch)
example['isRed'] = example.color.apply(RedMatch)
example['isOrange'] = example.color.apply(OrangeMatch)

prob  = pulp.LpProblem("complex_napsack", pulp.LpMaximize) 
x = pulp.LpVariable.dicts( "x", indexs = example.index, lowBound=0, upBound=1, cat='Integer', indexStart=[] )

prob += pulp.lpSum([x[i]*example.value[i] for i in example.index ])

prob += pulp.lpSum([x[i]*example.isBlue[i] for i in example.index])==1
prob += pulp.lpSum([x[i]*example.isYellow[i] for i in example.index])==1
prob += pulp.lpSum([x[i]*example.isRed[i] for i in example.index])==1
prob += pulp.lpSum([x[i]*example.isOrange[i] for i in example.index])==2

prob += pulp.lpSum([x[i] for i in example.index ])==5

prob.solve()

total_value = 0
for v in prob.variables():
    if v.varValue == 1.0:
        mystring = re.search('([0-9]*$)', v.name)
        print(v.name, "=", v.varValue)
        ind = int(mystring.group(1))
        print(example.item[ind])
        total_value = example.value[ind] + total_value

print(total_value)

#x_11 = 1.0
#L
#x_13 = 1.0
#N
#x_4 = 1.0
#E
#x_6 = 1.0
#G
#x_9 = 1.0
#J
#4.087

#But, the blue item (G) is a circle,  and both (L) and (N) are also circles.  I only want the red, orange, and yellow items to only have at most one shape in common with the blue item.

#How do I formulate a constraint to do this?

Thanks!

Answer 1

Several ways:

1) For each blue item, make sure that if it is selected, no more than 1 of that shape is selected from the non blue:

for idx in example.index:
  if example.isBlue[idx]:
    prob += pulp.lpSum([x[i] for i in example.index if not example.isBlue[i] and example.shapes[i] == example.shapes[idx]]) <= (1 - x[idx]) * 4 + 1

2) For each shape, make sure that if the blue was selected from that shape, no more than 1 of that shape is selected from the non blue:

for shape in example.shapes.unique():
  prob += pulp.lpSum([x[i] for i in example.index if not example.isBlue[i] and example.shapes[i] == shape]) <= (1 - pulp.lpSum([x[i] for i in example.index if example.isBlue[i] and example.shapes[i] == shape])) * 4 + 1  

I am making the right side be 1 when blue is selected and 5 if not (thus not imposing any restriction over the others).

3) Solve a problem for each shape, where you define the shape blue will have beforehand, make sure no more than 1 of non blue picks that shape, then stay with the best solution out of all problem solutions.