278 lines
9.3 KiB
Python
278 lines
9.3 KiB
Python
#!/usr/bin/env python3
|
|
"""budgeted_planner.py — exact optimal task selection under a cost budget and prerequisite closure.
|
|
|
|
Lane: Codeforge. Pure stdlib.
|
|
|
|
Problem
|
|
-------
|
|
Given tasks with a token `cost`, an integer `value`, and prerequisite ids `requires`,
|
|
choose a subset S that is
|
|
|
|
* CLOSED: every prerequisite of every task in S is also in S,
|
|
* FEASIBLE: sum(cost(t) for t in S) <= budget,
|
|
|
|
maximising sum(value(t) for t in S).
|
|
|
|
Exactness
|
|
---------
|
|
`topo_order` imposes an order where every prerequisite precedes its dependents.
|
|
A task's prerequisites are therefore all earlier in that order. The solver does a
|
|
memoised DP over (index, remaining budget, bitmask of *relevant* earlier selections).
|
|
A bit is "relevant" at index i iff it is the prerequisite of some task at index >= i;
|
|
restricting the mask to relevant bits is lossless because a non-relevant earlier
|
|
selection can never influence any future decision. No heuristic, no sampling: the DP
|
|
enumerates the whole decision space, so the returned value is the true optimum.
|
|
|
|
`brute_force` is an independent oracle (enumerate all 2**n subsets, test closure by set
|
|
containment, no topo order, no DP) used by the tests as the correctness reference.
|
|
|
|
Determinism
|
|
-----------
|
|
For a fixed input the output is byte-identical across runs. Among equally optimal
|
|
selections the solver deterministically prefers the one that includes a task when the
|
|
include/exclude values tie.
|
|
|
|
Limits
|
|
------
|
|
Exponential in the number of *relevant* earlier selections per step; comfortable for
|
|
tens of tasks. `brute_force` is 2**n and is only usable for n <= ~15.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from typing import Iterable, Mapping, Sequence
|
|
|
|
__all__ = [
|
|
"Task",
|
|
"SpecError",
|
|
"solve",
|
|
"brute_force",
|
|
"greedy",
|
|
"topo_order",
|
|
"check_selection",
|
|
"plan_from_json",
|
|
"main",
|
|
]
|
|
|
|
|
|
class SpecError(ValueError):
|
|
"""Raised for malformed input: duplicate ids, dangling/cyclic prerequisites,
|
|
non-positive costs, negative values, negative budget."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Task:
|
|
id: str
|
|
cost: int
|
|
value: int
|
|
requires: tuple[str, ...]
|
|
|
|
@classmethod
|
|
def of(cls, obj: Mapping) -> "Task":
|
|
return cls(
|
|
id=str(obj["id"]),
|
|
cost=int(obj["cost"]),
|
|
value=int(obj["value"]),
|
|
requires=tuple(str(r) for r in obj.get("requires", ())),
|
|
)
|
|
|
|
|
|
def _validate(tasks: Sequence[Task], budget: int) -> None:
|
|
if not isinstance(budget, int) or isinstance(budget, bool):
|
|
raise SpecError(f"budget must be an int, got {type(budget).__name__}")
|
|
if budget < 0:
|
|
raise SpecError(f"budget must be >= 0, got {budget}")
|
|
seen: set[str] = set()
|
|
for t in tasks:
|
|
if not t.id:
|
|
raise SpecError("task id must be non-empty")
|
|
if t.id in seen:
|
|
raise SpecError(f"duplicate task id: {t.id!r}")
|
|
seen.add(t.id)
|
|
if t.cost <= 0:
|
|
raise SpecError(f"task {t.id!r} cost must be > 0, got {t.cost}")
|
|
if t.value < 0:
|
|
raise SpecError(f"task {t.id!r} value must be >= 0, got {t.value}")
|
|
for t in tasks:
|
|
for r in t.requires:
|
|
if r not in seen:
|
|
raise SpecError(f"task {t.id!r} requires unknown task {r!r}")
|
|
|
|
|
|
def topo_order(ids: Iterable[str], requires: Mapping[str, Sequence[str]]) -> list[str]:
|
|
"""Return an order in which every prerequisite precedes its dependent.
|
|
|
|
Raises SpecError if the prerequisite graph has a cycle.
|
|
"""
|
|
pending = list(ids)
|
|
emitted: list[str] = []
|
|
done: set[str] = set()
|
|
while pending:
|
|
ready = sorted(
|
|
t for t in pending if all(r in done for r in requires.get(t, ()))
|
|
)
|
|
if not ready:
|
|
raise SpecError(
|
|
"prerequisite cycle among: " + ", ".join(sorted(pending))
|
|
)
|
|
for t in ready:
|
|
emitted.append(t)
|
|
done.add(t)
|
|
pending.remove(t)
|
|
return emitted
|
|
|
|
|
|
def _relevant_masks(order: Sequence[str], req_idx: Mapping[str, tuple[int, ...]]) -> list[frozenset[int]]:
|
|
"""relevant[i] = indices < i that are a prerequisite of some task at index >= i."""
|
|
n = len(order)
|
|
rel: list[frozenset[int]] = []
|
|
for i in range(n + 1):
|
|
acc: set[int] = set()
|
|
for j in range(i, n):
|
|
acc.update(req_idx[order[j]])
|
|
rel.append(frozenset(j for j in acc if j < i))
|
|
return rel
|
|
|
|
|
|
def solve(tasks: Sequence[Task], budget: int) -> dict:
|
|
"""Exact optimum. Returns {'selected', 'cost', 'value'}."""
|
|
tasks = list(tasks)
|
|
_validate(tasks, budget)
|
|
order = topo_order([t.id for t in tasks], {t.id: t.requires for t in tasks})
|
|
idx = {t: i for i, t in enumerate(order)}
|
|
cost = {t.id: t.cost for t in tasks}
|
|
value = {t.id: t.value for t in tasks}
|
|
req_idx = {t.id: tuple(sorted(idx[r] for r in t.requires)) for t in tasks}
|
|
rel = _relevant_masks(order, req_idx)
|
|
n = len(order)
|
|
|
|
memo: dict[tuple[int, int, int], int] = {}
|
|
|
|
def f(i: int, rem: int, mask: int) -> int:
|
|
if i == n:
|
|
return 0
|
|
key = (i, rem, mask)
|
|
hit = memo.get(key)
|
|
if hit is not None:
|
|
return hit
|
|
best = f(i + 1, rem, mask & _bits(rel[i + 1]))
|
|
t = order[i]
|
|
if cost[t] <= rem and all(mask >> r & 1 for r in req_idx[t]):
|
|
cand = value[t] + f(i + 1, rem - cost[t], (mask | (1 << i)) & _bits(rel[i + 1]))
|
|
if cand > best:
|
|
best = cand
|
|
memo[key] = best
|
|
return best
|
|
|
|
chosen: list[str] = []
|
|
|
|
def build(i: int, rem: int, mask: int) -> None:
|
|
if i == n:
|
|
return
|
|
t = order[i]
|
|
excl = f(i + 1, rem, mask & _bits(rel[i + 1]))
|
|
if cost[t] <= rem and all(mask >> r & 1 for r in req_idx[t]):
|
|
inc = value[t] + f(i + 1, rem - cost[t], (mask | (1 << i)) & _bits(rel[i + 1]))
|
|
if inc >= excl: # deterministic tie-break: include
|
|
chosen.append(t)
|
|
build(i + 1, rem - cost[t], (mask | (1 << i)) & _bits(rel[i + 1]))
|
|
return
|
|
build(i + 1, rem, mask & _bits(rel[i + 1]))
|
|
|
|
build(0, budget, 0)
|
|
out = sorted(chosen)
|
|
result = {
|
|
"selected": out,
|
|
"cost": sum(cost[t] for t in out),
|
|
"value": sum(value[t] for t in out),
|
|
}
|
|
if result["value"] != f(0, budget, 0): # internal consistency assertion
|
|
raise AssertionError("reconstruction diverged from DP value")
|
|
return result
|
|
|
|
|
|
def _bits(fs: frozenset[int]) -> int:
|
|
b = 0
|
|
for j in fs:
|
|
b |= 1 << j
|
|
return b
|
|
|
|
|
|
def brute_force(tasks: Sequence[Task], budget: int) -> int:
|
|
"""Independent oracle: enumerate every subset, keep closed+feasible ones, return max value."""
|
|
tasks = list(tasks)
|
|
_validate(tasks, budget)
|
|
ids = [t.id for t in tasks]
|
|
cost = {t.id: t.cost for t in tasks}
|
|
value = {t.id: t.value for t in tasks}
|
|
req = {t.id: frozenset(t.requires) for t in tasks}
|
|
best = 0
|
|
for mask in range(1 << len(ids)):
|
|
chosen = {ids[i] for i in range(len(ids)) if mask >> i & 1}
|
|
if not chosen:
|
|
continue
|
|
if any(not req[i] <= chosen for i in chosen):
|
|
continue
|
|
if sum(cost[i] for i in chosen) > budget:
|
|
continue
|
|
v = sum(value[i] for i in chosen)
|
|
if v > best:
|
|
best = v
|
|
return best
|
|
|
|
|
|
def greedy(tasks: Sequence[Task], budget: int) -> dict:
|
|
"""Cheap lower bound: scan in id order, take any task whose prerequisites are
|
|
all already taken and whose cost fits. Not exact; used as a sanity floor."""
|
|
tasks = list(tasks)
|
|
_validate(tasks, budget)
|
|
order = topo_order([t.id for t in tasks], {t.id: t.requires for t in tasks})
|
|
by_id = {t.id: t for t in tasks}
|
|
taken: set[str] = set()
|
|
spent = 0
|
|
gained = 0
|
|
for tid in order:
|
|
t = by_id[tid]
|
|
if all(r in taken for r in t.requires) and spent + t.cost <= budget:
|
|
taken.add(tid)
|
|
spent += t.cost
|
|
gained += t.value
|
|
return {"selected": sorted(taken), "cost": spent, "value": gained}
|
|
|
|
|
|
def check_selection(tasks: Sequence[Task], budget: int, selected: Sequence[str]) -> None:
|
|
"""Raise AssertionError unless `selected` is closed, feasible and internally consistent."""
|
|
by_id = {t.id: t for t in tasks}
|
|
s = set(selected)
|
|
assert len(s) == len(selected), "selected contains duplicates"
|
|
assert s <= set(by_id), f"unknown task(s): {sorted(s - set(by_id))}"
|
|
for t in s:
|
|
missing = set(by_id[t].requires) - s
|
|
assert not missing, f"task {t!r} is missing prerequisites {sorted(missing)}"
|
|
assert sum(by_id[t].cost for t in s) <= budget, "selection exceeds budget"
|
|
|
|
|
|
def plan_from_json(payload: Mapping) -> dict:
|
|
tasks = [Task.of(o) for o in payload.get("tasks", ())]
|
|
return solve(tasks, int(payload.get("budget", 0)))
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
|
argv = list(sys.argv[1:] if argv is None else argv)
|
|
raw = open(argv[0], "r", encoding="utf-8").read() if argv else sys.stdin.read()
|
|
try:
|
|
payload = json.loads(raw)
|
|
result = plan_from_json(payload)
|
|
except (json.JSONDecodeError, SpecError, KeyError) as exc:
|
|
print(json.dumps({"error": f"{type(exc).__name__}: {exc}"}), file=sys.stderr)
|
|
return 2
|
|
print(json.dumps(result, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|