time to splat combinations calc

splatoon_ttk.py

from itertools import combinations_with_replacement
from collections import namedtuple

Attack = namedtuple('Attack', ['damage', 'label'])

TARGET_DAMAGE = 99
"""If the combo has total damage above this, it is recorded"""
MAX_ATTACKS = 3
"""The maximum number of attack attempts to hit the target"""

# Translation table for damage values
damage_table = {
    "main": [
        Attack(33, "main direct"),
        Attack(28, "main off-hit"),
    ],
    "main2": [
        Attack(36, "main direct"),
        Attack(32, "main off-hit"),
    ],
    "sub": [
        Attack(80, "sub direct"),
        Attack(30, "sub off-hit"),
    ],
    "sub2": [
        Attack(60, "sub direct"),
        Attack(20, "sub off-hit"),
    ],
    "special": [
        Attack(120, "special direct"),
        Attack(60, "special off-hit"),
    ],
    "special2": [
        Attack(99, "special direct"),
        Attack(80, "special off-hit"),
    ]
}


def find_combos(main: str, sub: str, special: str) -> tuple[Attack]:
    """For the given set of main, sub, and special, find the combos that are TARGET_DAMAGE or above within MAX_ATTACKS."""
    attacks = damage_table[main] + damage_table[sub] + damage_table[special]
    attacks = sorted(attacks, key=lambda a: a.damage, reverse=True)
    results = set()
    for combo in combinations_with_replacement(attacks, r=MAX_ATTACKS):
        total_damage = 0
        for i, attack in enumerate(combo):
            total_damage += attack.damage
            if total_damage >= TARGET_DAMAGE:
                if add_set_success(results, combo[:i+1]):
                    yield combo[:i+1]
                break


def add_set_success(s, item) -> bool:
    return len(s) != (s.add(item) or len(s))


def ask_for_arg(label: str):
    while True:
        temp = input("Enter the " + label)
        if temp in damage_table:
            return temp
        else:
            print(temp + " not found.")


if __name__ == "__main__":
    _main = ask_for_arg("main")
    _sub = ask_for_arg("sub")
    _special = ask_for_arg("special")
    print("Combos that result in {}+ damage:".format(TARGET_DAMAGE))
    for _combo in find_combos(_main, _sub, _special):
        print(_combo)