Replace brute-force laptop allocation with scalable greedy approach; Remove unwanted imports

Author: ike-agu

sprint-5-laptop-allocation/laptop_allocation.py

@@ -1,14 +1,12 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import List, Dict, Optional, Tuple
-import itertools
+from typing import List, Dict, Optional
 import sys
 
 class OperatingSystem(Enum):
     MACOS = "macOS"
     ARCH = "Arch Linux"
     UBUNTU = "Ubuntu"
-
 @dataclass(frozen=True)
 class Person:
     name: str
@@ -18,8 +16,6 @@ class Person:
 # custom hash to hash the immutable version of the list
     def __hash__(self) -> int:
         return hash((self.name, self.age, tuple(self.preferred_operating_system)))
-
-
 @dataclass(frozen=True)
 class Laptop:
     id: int
@@ -30,22 +26,23 @@ class Laptop:
 
 # =====define parse operating system======
 def parse_operating_system(raw: str) -> OperatingSystem:
-    normalized = raw.strip().lower()
+   normalized = raw.strip().lower()
 
-    aliases = {
-        "macos": OperatingSystem.MACOS,
-        "mac": OperatingSystem.MACOS,
-        "arch": OperatingSystem.ARCH,
-        "arch linux": OperatingSystem.ARCH,
-        "ubuntu": OperatingSystem.UBUNTU,
-    }
+   aliases = {
+       "macos": OperatingSystem.MACOS,
+       "mac": OperatingSystem.MACOS,
+       "arch": OperatingSystem.ARCH,
+       "arch linux": OperatingSystem.ARCH,
+       "ubuntu": OperatingSystem.UBUNTU,
+   }
 
-    if normalized not in aliases:
-        valid = ", ".join(sorted(aliases.keys()))
-        print(f"Error: invalid operating system. Try one of: {valid}", file=sys.stderr)
-        sys.exit(1)
 
-    return aliases[normalized]
+   if normalized not in aliases:
+       valid = ", ".join(sorted(aliases.keys()))
+       print(f"Error: invalid operating system. Try one of: {valid}", file=sys.stderr)
+       sys.exit(1)
+
+   return aliases[normalized]
 
 # ======= sadness helper ============
 def sadness(person: Person, laptop: Laptop) -> int:
@@ -60,25 +57,30 @@ def allocate_laptops(people: List[Person], laptops: List[Laptop]) -> Dict[Person
     if len(laptops) < len(people):
         raise ValueError("Not enough laptops to allocate one per person.")
 
-    best_assignment: Optional[Dict[Person, Laptop]] = None
-    best_total_sadness: Optional[int] = None
+    sorted_people = sorted(people, key=lambda p: len(p.preferred_operating_system))
+
+    remaining_laptops = laptops[:]
+    assignment: Dict[Person, Laptop] = {}
+
+    for person in sorted_people:
+        best_laptop:Optional[Laptop] = None
+        best_score = 10**9
 
-    for chosen_laptops in itertools.combinations(laptops, len(people)):
-        for permuted_laptops in itertools.permutations(chosen_laptops):
-            total = 0
-            for i in range(len(people)):
-                person = people[i]
-                laptop = permuted_laptops[i]
-                total += sadness(person, laptop)
+        for laptop in remaining_laptops:
+            score = sadness(person, laptop)
+            if score < best_score:
+                best_score = score
+                best_laptop = laptop
+                if score == 0:
+                    break
 
-            if best_total_sadness is None or total < best_total_sadness:
-                best_total_sadness = total
-                best_assignment = {people[i]: permuted_laptops[i] for i in range(len(people))}
+        if best_laptop is None:
+            raise RuntimeError("Allocation failed unexpectedly.")
 
-    if best_assignment is None:
-        raise RuntimeError("Allocation failed unexpectedly.")
+        assignment[person] = best_laptop
+        remaining_laptops.remove(best_laptop)
 
-    return best_assignment
+    return assignment
 
 # ======define main ================
 def main() -> None: