from time import perf_counter

class ExecutionTimer:
    def __init__(self, message: str):
        self.message = message

    def __enter__(self):
        self.start = perf_counter()
        return self

    def __exit__(self, type, value, traceback):
        self.time = perf_counter() - self.start
        print(f'{self.message}: {self.time:.3f} seconds')

import pulp, numpy as np, itertools

# Some reward matrix
reward = np.full((300, 200, 50), 10)

prob = pulp.LpProblem("Problem", pulp.constants.LpMaximize)
dv = pulp.LpVariable.dicts("dv",
       itertools.product(range(300), range(200), range(50)),
       cat=pulp.constants.LpBinary)

with ExecutionTimer("Slow Objective") as _:
    prob += pulp.lpSum(
        dv[a, b, c] * reward[a, b, c]
        for a in range(300)
        for b in range(200)
        for c in range(50)
    )

with ExecutionTimer("Slow Constraint") as _:
    for c in range(50):
        prob += pulp.lpSum(
            dv[a, b, c]
            for a in range(300)
            for b in range(200)
        ) == 1

prob = pulp.LpProblem("Problem", pulp.constants.LpMaximize)
dv = pulp.LpVariable.dicts("dv",
       itertools.product(range(300), range(200), range(50)),
       cat=pulp.constants.LpBinary)

with ExecutionTimer("Fast Objective") as _:
    prob += pulp.LpAffineExpression(
        (dv[a, b, c], reward[a, b, c])
        for a in range(300)
        for b in range(200)
        for c in range(50)
    )

with ExecutionTimer("Fast Constraint") as _:
    for c in range(50):
        prob += pulp.LpAffineExpression(
            (dv[a, b, c], 1)
            for a in range(300)
            for b in range(200)
        ) == 1


def __eq__(self, other):
    if isinstance(other, (float, int)):
        return pulp.LpConstraint(self, pulp.const.LpConstraintEQ, rhs=other)
    else:
        return pulp.LpConstraint(self - other, pulp.const.LpConstraintEQ)

pulp.LpAffineExpression.__eq__ = __eq__

with ExecutionTimer("Faster Constraint") as _:
    for c in range(50):
        prob += pulp.LpAffineExpression(
            (dv[a, b, c], 1)
            for a in range(300)
            for b in range(200)
        ) == 1