95 lines
2.8 KiB
Python
95 lines
2.8 KiB
Python
""" Implements Multiway number partitioning greedy algorithm """
|
|
|
|
import typing as t
|
|
import logging
|
|
from sortedcontainers import SortedList
|
|
from .config import Task
|
|
|
|
__all__ = ["TaskId", "partition"]
|
|
|
|
TaskId = t.NewType("TaskId", int)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class PartitionException(Exception):
|
|
"""An exception occurring during partitioning"""
|
|
|
|
|
|
class UnsolvableConflict(PartitionException):
|
|
"""Cannot partition set due to unsolvable conflicts"""
|
|
|
|
|
|
class Bin:
|
|
"""A bin containing assigned tasks"""
|
|
|
|
elts: list[TaskId]
|
|
cost: int
|
|
|
|
def __init__(self):
|
|
self.elts = []
|
|
self.cost = 0
|
|
|
|
def add(self, task: TaskId, cost: int):
|
|
assert task not in self.elts
|
|
self.elts.append(task)
|
|
self.cost += cost
|
|
|
|
def __contains__(self, task: TaskId) -> bool:
|
|
return task in self.elts
|
|
|
|
|
|
def partition(
|
|
bin_count: int,
|
|
tasks: list[Task],
|
|
costs: dict[TaskId, int],
|
|
multiplicity: dict[TaskId, int],
|
|
) -> list[list[TaskId]]:
|
|
"""Partitions the tasks, each with cost `costs[i]`, into `bin_count` bins. Each
|
|
task has multiplicity `multiplicity[i]`, copies of the same task being mutually
|
|
exclusive (ie. cannot be in the same bin)"""
|
|
|
|
bins = SortedList([Bin() for _ in range(bin_count)], key=lambda x: x.cost)
|
|
ordered_tasks: list[TaskId] = []
|
|
for t_id, reps in multiplicity.items():
|
|
for _ in range(reps):
|
|
ordered_tasks.append(t_id)
|
|
ordered_tasks.sort(key=lambda x: costs[x], reverse=True)
|
|
|
|
for task in ordered_tasks:
|
|
possible_bins: list[int] = []
|
|
min_possible: t.Optional[int] = None
|
|
for pos, cur_bin in enumerate(bins):
|
|
if min_possible is not None and cur_bin.cost > min_possible:
|
|
break
|
|
if task not in cur_bin:
|
|
if min_possible is None:
|
|
min_possible = cur_bin.cost
|
|
possible_bins.append(pos)
|
|
if not possible_bins:
|
|
raise UnsolvableConflict(
|
|
"Pas assez de groupes pour affecter la tâche "
|
|
+ f"{tasks[task].qualified_name} {multiplicity[task]} fois."
|
|
)
|
|
|
|
# Pick one of the groups -- maximize distance to other tasks
|
|
closest_assigned = []
|
|
for cur_bin_id in possible_bins:
|
|
cur_bin = bins[cur_bin_id]
|
|
if cur_bin.elts:
|
|
closest_assigned.append(min((abs(elt - task) for elt in cur_bin.elts)))
|
|
else:
|
|
closest_assigned.append(len(costs) + 1)
|
|
|
|
assign_to_bin_id, _ = max(enumerate(closest_assigned), key=lambda x: x[1])
|
|
assign_to_bin = bins[possible_bins[assign_to_bin_id]]
|
|
|
|
del bins[assign_to_bin_id]
|
|
assign_to_bin.add(task, costs[task])
|
|
bins.add(assign_to_bin)
|
|
|
|
out: list[list[TaskId]] = []
|
|
for cur_bin in bins:
|
|
out.append(cur_bin.elts)
|
|
return out
|