結果

問題 No.2869 yuusaan's Knapsacks
ユーザー hitonanodehitonanode
提出日時 2024-08-31 00:17:14
言語 Python3
(3.12.2 + numpy 1.26.4 + scipy 1.12.0)
結果
AC  
実行時間 1,794 ms / 4,500 ms
コード長 11,764 bytes
コンパイル時間 290 ms
コンパイル使用メモリ 13,824 KB
実行使用メモリ 71,208 KB
最終ジャッジ日時 2024-08-31 00:18:05
合計ジャッジ時間 45,966 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge1
外部呼び出し有り
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 1,369 ms
70,356 KB
testcase_01 AC 1,374 ms
70,096 KB
testcase_02 AC 1,373 ms
70,868 KB
testcase_03 AC 1,375 ms
70,612 KB
testcase_04 AC 1,417 ms
70,356 KB
testcase_05 AC 1,365 ms
69,556 KB
testcase_06 AC 1,378 ms
70,424 KB
testcase_07 AC 1,574 ms
70,936 KB
testcase_08 AC 1,371 ms
70,364 KB
testcase_09 AC 1,352 ms
69,576 KB
testcase_10 AC 1,492 ms
70,660 KB
testcase_11 AC 1,404 ms
70,384 KB
testcase_12 AC 1,424 ms
71,208 KB
testcase_13 AC 1,392 ms
70,776 KB
testcase_14 AC 1,794 ms
71,180 KB
testcase_15 AC 1,406 ms
70,348 KB
testcase_16 AC 1,384 ms
70,488 KB
testcase_17 AC 1,378 ms
70,732 KB
testcase_18 AC 1,426 ms
70,344 KB
testcase_19 AC 1,485 ms
70,904 KB
testcase_20 AC 1,386 ms
69,960 KB
testcase_21 AC 1,429 ms
70,352 KB
testcase_22 AC 1,392 ms
70,872 KB
testcase_23 AC 1,405 ms
70,600 KB
testcase_24 AC 1,419 ms
70,480 KB
testcase_25 AC 1,400 ms
70,124 KB
testcase_26 AC 1,390 ms
70,084 KB
testcase_27 AC 1,378 ms
69,836 KB
testcase_28 AC 1,416 ms
69,960 KB
testcase_29 AC 1,400 ms
70,556 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

from enum import Enum
from typing import SupportsFloat

from scipy.optimize import linprog  # type: ignore


class LPVarType(Enum):
    CONTINUOUS = 0
    INTEGER = 1
    SEMICONTINUOUS = 2  # 0 or lb <= x <= ub
    SEMIINTEGER = 3  # 0 or lb <= x <= ub, integer


class LPStatus(Enum):
    OPTIMAL = 0
    ITERATION_LIMIT = 1
    INFEASIBLE = 2
    UNBOUNDED = 3
    NUMERICAL_ERROR = 4


class LPSense(Enum):
    MINIMIZE = 0
    MAXIMIZE = 1


class _LPBase:
    def __add__(self, other: "LPExpressionLike") -> "LPExpression":
        return LPExpression.build(self) + other

    def __radd__(self, other: "LPExpressionLike") -> "LPExpression":
        return other + LPExpression.build(self)

    def __sub__(self, other: "LPExpressionLike") -> "LPExpression":
        return LPExpression.build(self) - other

    def __rsub__(self, other: "LPExpressionLike") -> "LPExpression":
        return other - LPExpression.build(self)

    def __mul__(self, other: SupportsFloat) -> "LPExpression":
        return LPExpression.build(self) * other

    def __rmul__(self, other: SupportsFloat) -> "LPExpression":
        return LPExpression.build(self) * other

    def __truediv__(self, other: SupportsFloat) -> "LPExpression":
        return LPExpression.build(self) / other

    def __neg__(self) -> "LPExpression":
        return -LPExpression.build(self)

    def __le__(
        self,
        other: "LPExpressionLike",
    ) -> "LPInequality":
        return LPInequality(
            lhs=LPExpression.build(self),
            rhs=other,
            inequality_type=LPInequalityType.LESSEQ,
        )

    def __ge__(
        self,
        other: "LPExpressionLike",
    ) -> "LPInequality":
        return LPInequality(
            lhs=other,
            rhs=LPExpression.build(self),
            inequality_type=LPInequalityType.LESSEQ,
        )

    def __eq__(  # type: ignore
        self,
        other: "LPExpressionLike",
    ) -> "LPInequality":
        return LPInequality(
            lhs=LPExpression.build(self),
            rhs=other,
            inequality_type=LPInequalityType.EQUAL,
        )


class LPVar(_LPBase):
    def __init__(
        self,
        name: str,
        lower_bound: float | None = None,
        upper_bound: float | None = None,
        variable_type: LPVarType = LPVarType.CONTINUOUS,
    ) -> None:
        self.name = name
        self.lower_bound: float | None = lower_bound
        self.upper_bound: float | None = upper_bound
        self.variable_type: LPVarType = variable_type
        self._value: float | None = None

    def __str__(self) -> str:
        s = "{}(lb={}, ub={}, type={})".format(
            self.name,
            self.lower_bound,
            self.upper_bound,
            self.variable_type.name,
        )
        if self._value is not None:
            s += ": {}".format(self._value)
        return s

    def __repr__(self) -> str:
        return str(self)

    def value(self) -> float | None:
        return self._value


class LPInequalityType(Enum):
    LESSEQ = 1  # lhs <= rhs
    EQUAL = 2  # lhs == rhs


class _LPTerm(_LPBase):
    """
    A term of the form coefficient * variable.
    """

    def __init__(self, coefficient: float, variable: LPVar) -> None:
        self.coefficient: float = coefficient
        self.variable: LPVar = variable


class LPExpression:
    def __init__(
        self,
        const: SupportsFloat,
        terms: list[_LPTerm],
    ) -> None:
        self.const = float(const)
        self.terms = terms.copy()
        self._value: float | None = None

    @classmethod
    def build(self, x: "LPExpressionLike | _LPBase") -> "LPExpression":
        if isinstance(x, LPExpression):
            return LPExpression(x.const, x.terms)
        elif isinstance(x, _LPTerm):
            return LPExpression(0, [x])
        elif isinstance(x, LPVar):
            return LPExpression(0, [_LPTerm(1.0, x)])
        elif isinstance(x, SupportsFloat):
            return LPExpression(x, [])
        else:
            raise TypeError("Invalid type for LPExpression")

    def __add__(self, other: "LPExpressionLike") -> "LPExpression":
        rhs = LPExpression.build(other)
        return LPExpression(
            self.const + rhs.const,
            self.terms + rhs.terms,
        )

    def __radd__(self, other: "LPExpressionLike") -> "LPExpression":
        return self + other

    def __sub__(self, other: "LPExpressionLike") -> "LPExpression":
        rhs = LPExpression.build(other)
        return self + (-rhs)

    def __rsub__(self, other: "LPExpressionLike") -> "LPExpression":
        return LPExpression.build(other) + (-self)

    def __mul__(self, other: SupportsFloat) -> "LPExpression":
        return LPExpression(
            self.const * float(other),
            [
                _LPTerm(
                    t.coefficient * float(other),
                    t.variable,
                )
                for t in self.terms
            ],
        )

    def __rmul__(self, other: SupportsFloat) -> "LPExpression":
        return self * other

    def __truediv__(self, other: SupportsFloat) -> "LPExpression":
        return LPExpression(
            self.const / float(other),
            [
                _LPTerm(
                    t.coefficient / float(other),
                    t.variable,
                )
                for t in self.terms
            ],
        )

    def __neg__(self) -> "LPExpression":
        return LPExpression(
            -self.const,
            [_LPTerm(-t.coefficient, t.variable) for t in self.terms],
        )

    def __le__(
        self,
        other: "LPExpressionLike",
    ) -> "LPInequality":
        return LPInequality(self, other, LPInequalityType.LESSEQ)

    def __ge__(
        self,
        other: "LPExpressionLike",
    ) -> "LPInequality":
        return LPInequality(other, self, LPInequalityType.LESSEQ)

    def __eq__(  # type: ignore
        self,
        other: "LPExpressionLike",
    ) -> "LPInequality":
        return LPInequality(self, other, LPInequalityType.EQUAL)

    def __str__(self) -> str:
        ret = ["{}".format(self.const)]

        for t in self.terms:
            sgn = "+" if t.coefficient >= 0 else ""
            ret.append("{}{}{}".format(sgn, t.coefficient, t.variable.name))

        return " ".join(ret)

    def value(self) -> float | None:
        return self._value


LPExpressionLike = LPExpression | _LPTerm | LPVar | SupportsFloat


class LPInequality:
    def __init__(
        self,
        lhs: LPExpressionLike,
        rhs: LPExpressionLike,
        inequality_type: LPInequalityType,
    ) -> None:
        """
        lhs <= rhs
        -> terms + const (<= or ==) 0
        """

        self.lhs = LPExpression.build(lhs) - LPExpression.build(rhs)
        self.inequality_type = inequality_type

    def __str__(self) -> str:
        if self.inequality_type == LPInequalityType.LESSEQ:
            return "{} <= 0".format(self.lhs)
        elif self.inequality_type == LPInequalityType.EQUAL:
            return "{} == 0".format(self.lhs)
        else:
            raise ValueError("Invalid inequality type")

    def __repr__(self) -> str:
        return str(self)


class LPModel:
    def __init__(self, sense: LPSense = LPSense.MINIMIZE) -> None:
        self.sense = sense
        self.has_impossible_constraints = False
        self.constraints: list[LPInequality] = []
        self.objective: LPExpression = LPExpression(0, [])
        self.status: LPStatus | None = None

    def add_constraint(self, constraint: LPInequality | bool) -> None:
        if isinstance(constraint, bool):
            if not constraint:
                self.has_impossible_constraints = True
        else:
            self.constraints.append(constraint)

    def set_objective(self, objective: LPExpressionLike) -> None:
        self.objective = LPExpression.build(objective)

    def solve(self) -> None:
        var_dict: dict[int, LPVar] = {}
        for constraint in self.constraints:
            for term in constraint.lhs.terms:
                var_dict.setdefault(id(term.variable), term.variable)

        for term in self.objective.terms:
            var_dict.setdefault(id(term.variable), term.variable)

        # Reset status
        self.objective._value = None
        self.status = None
        for var in var_dict.values():
            var._value = None

        # Obviously infeasible
        if self.has_impossible_constraints:
            self.status = LPStatus.INFEASIBLE
            return

        # Obviously optimal
        if not var_dict:
            self.status = LPStatus.OPTIMAL
            self.objective._value = self.objective.const
            return

        id_to_idx = {id(v): i for i, v in enumerate(var_dict.values())}

        A_ub: list[list[float]] = []
        b_ub: list[float] = []
        A_eq: list[list[float]] = []
        b_eq: list[float] = []

        for constraint in self.constraints:
            lhs: list[float] = [0.0] * len(var_dict)
            rhs = -constraint.lhs.const

            for term in constraint.lhs.terms:
                lhs[id_to_idx[id(term.variable)]] += term.coefficient

            if constraint.inequality_type == LPInequalityType.LESSEQ:
                A_ub.append(lhs)
                b_ub.append(rhs)
            elif constraint.inequality_type == LPInequalityType.EQUAL:
                A_eq.append(lhs)
                b_eq.append(rhs)
            else:
                raise ValueError("Invalid inequality type")

        bounds = [(v.lower_bound, v.upper_bound) for v in var_dict.values()]

        integrality = [v.variable_type.value for v in var_dict.values()]

        c: list[float] = [0.0] * len(var_dict)

        func_weight = 1 if self.sense == LPSense.MINIMIZE else -1

        for term in self.objective.terms:
            c[id_to_idx[id(term.variable)]] += term.coefficient * func_weight

        res = linprog(
            c,
            A_ub=A_ub or None,
            b_ub=b_ub or None,
            A_eq=A_eq or None,
            b_eq=b_eq or None,
            bounds=bounds,
            integrality=integrality if sum(integrality) else None,
        )

        if res.status == 0:
            self.status = LPStatus.OPTIMAL

            for i, variable in enumerate(var_dict.values()):
                variable._value = res.x[i]

            self.objective._value = res.fun * func_weight + self.objective.const

        elif res.status == 1:
            self.status = LPStatus.ITERATION_LIMIT

        elif res.status == 2:
            self.status = LPStatus.INFEASIBLE

        elif res.status == 3:
            self.status = LPStatus.UNBOUNDED

        elif res.status == 4:
            self.status = LPStatus.NUMERICAL_ERROR

        else:
            raise ValueError("Invalid status code")


n, m = map(int, input().split())
e = list(map(int, input().split()))

v: list[int] = []
w: list[int] = []
for j in range(m):
    a, b = map(int, input().split())
    v.append(a)
    w.append(b)

model = LPModel(sense=LPSense.MAXIMIZE)

var = [
    [
        LPVar("x{}_{}".format(i, j), variable_type=LPVarType.INTEGER, lower_bound=0, upper_bound=1)
        for j in range(m)
    ]
    for i in range(n)  # 箱 i に物 j
]

for j in range(m):
    model.add_constraint(sum(var[i][j] for i in range(n)) <= 1)

for i in range(n):
    model.add_constraint(sum(var[i][j] * w[j] for j in range(m)) <= e[i])

model.set_objective(sum(var[i][j] * v[j] for i in range(n) for j in range(m)))

model.solve()

print(round(model.objective.value()))

ret: list[list[int]] = [[] for _ in range(n)]

for i in range(n):
    for j in range(m):
        if var[i][j].value() > 0.9:
            ret[i].append(j + 1)

for r in ret:
    print(len(r), *r)
0