refactor: improve knapsack selection with dynamic scaling, pre-filtering, and greedy budget filling

oh_security
Cizz22 3 months ago
parent d2ecf0849e
commit 19066475c9

@ -1,6 +1,7 @@
from collections import defaultdict
import math
import random
from typing import Optional
from typing import Optional, List, Dict, Tuple
from uuid import UUID
from sqlalchemy import Delete, Select
@ -112,61 +113,88 @@ def greedy_selection(equipments: List[dict], budget: float) -> Tuple[List[dict],
"""Greedy fallback: select items by score until budget is used."""
# Sort by priority_score descending
equipments_sorted = sorted(equipments, key=lambda x: x["priority_score"], reverse=True)
total_cost = 0
current_cost = 0.0
selected, excluded = [], []
for eq in equipments_sorted:
if total_cost + eq["cost"] <= budget:
cost = eq.get("total_cost", 0.0)
if current_cost + cost <= budget:
selected.append(eq)
total_cost += eq["cost"]
current_cost += cost
else:
excluded.append(eq)
return selected, excluded
def knapsack_selection(equipments: List[dict], budget: float, scale: int = 10_000_000) -> Tuple[List[dict], List[dict]]:
def knapsack_selection(equipments: List[dict], budget: float, scale: Optional[float] = None) -> Tuple[List[dict], List[dict]]:
"""
Select equipment optimally within budget using 0/1 knapsack DP.
Uses scaling + 1D DP optimization to avoid MemoryError.
Falls back to greedy if W is too large.
Uses dynamic scaling + 1D DP optimization to ensure accuracy and avoid MemoryError.
"""
n = len(equipments)
if not equipments:
return [], []
# Filter out items that are strictly above budget right away
# This helps significantly when budget is very low (e.g. 1)
eligible_items = []
strictly_excluded = []
for eq in equipments:
if eq["total_cost"] > budget:
strictly_excluded.append(eq)
else:
eligible_items.append(eq)
if not eligible_items:
return [], strictly_excluded
# Scale costs + budget
costs = [int(eq["total_cost"] // scale) for eq in equipments]
values = [eq["priority_score"] for eq in equipments]
n = len(eligible_items)
# Dynamic scaling: target W around 2000 for good precision/performance balance
if scale is None:
target_W = 2000
scale = max(1.0, budget / target_W)
costs = [int(math.ceil(eq["total_cost"] / scale)) for eq in eligible_items]
values = [eq["priority_score"] for eq in eligible_items]
W = int(budget // scale)
# Fallback if W is still too large
# Fallback if W is still too large (memory safety)
if W > 1_000_000:
print("too big")
return greedy_selection(equipments, budget)
# 1D DP array
dp = [0.0] * (W + 1)
keep = [[False] * (W + 1) for _ in range(n)] # track selection choices
keep = [[False] * (W + 1) for _ in range(n)]
for i in range(n):
cost, value = costs[i], values[i]
for w in range(W, cost - 1, -1):
if dp[w - cost] + value >= dp[w]: # <= FIXED HERE
if dp[w - cost] + value >= dp[w]:
dp[w] = dp[w - cost] + value
keep[i][w] = True
# Backtrack to find selected items
selected, excluded = [], []
# Backtrack
selected = []
backtrack_excluded = []
w = W
for i in range(n - 1, -1, -1):
if keep[i][w]:
selected.append(equipments[i])
selected.append(eligible_items[i])
w -= costs[i]
else:
excluded.append(equipments[i])
backtrack_excluded.append(eligible_items[i])
excluded = backtrack_excluded + strictly_excluded
# Optional: greedy fill leftover budget with actual values
# This compensates for any conservative rounding from math.ceil
current_total_cost = sum(eq["total_cost"] for eq in selected)
remaining_budget = budget - current_total_cost
# Optional: fill leftover budget with zero-priority items
remaining_budget = budget - sum(eq["total_cost"] for eq in selected)
if remaining_budget > 0:
# Sort excluded items by priority score for better greedy fill
excluded.sort(key=lambda x: x["priority_score"], reverse=True)
for eq in excluded[:]:
if eq["total_cost"] <= remaining_budget:
selected.append(eq)

Loading…
Cancel
Save