diff --git a/src/calculation_budget_constrains/service.py b/src/calculation_budget_constrains/service.py index d16d7b5..4a838b4 100644 --- a/src/calculation_budget_constrains/service.py +++ b/src/calculation_budget_constrains/service.py @@ -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]) - # Optional: fill leftover budget with zero-priority items - remaining_budget = budget - sum(eq["total_cost"] for eq in selected) + 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 + 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)