Minimum-cost encryption of a directory tree

Every file in the tree must end up encrypted. You have two calls:
encrypt(file) costs file_cost, and encrypt(directory) costs dir_cost
and encrypts every file anywhere below that directory. Re-encrypting a file
that is already encrypted is harmless in the base version. Find the cheapest
set of calls.

Part 1: count first. A plain DFS over directories and files gives the
number of encrypted and unencrypted files. Use an explicit stack: real
filesystems can be deeper than Python's recursion limit.

Brute force. Choose, for every directory, whether to call encrypt on it,
and add a file call for every unencrypted file not covered. That is 2^D
choices for D directories. Too slow, but it shows the structure. Once you
encrypt a directory, nothing inside it matters. So the choices for one subtree
do not affect its siblings.

Optimal: tree DP, post-order. For a directory d, let best(d) be the
cheapest cost to encrypt every file in its subtree:

children(d) = file_cost * (unencrypted direct files) + sum(best(c) for c in subdirs)
best(d)     = 0                          if no unencrypted file under d
            = min(children(d), dir_cost) otherwise

This is correct because a call on an ancestor already covers d. So inside
the subproblem for d, the only question is whether to cover d with one
call or to pay for its parts. The parts are independent subproblems. Compute
children before parents (post-order) and carry the call list with the cost so
you can return the plan, not just the number.

Invariants:

• best(d) <= children(d), and best(d) <= dir_cost whenever there is work.
• A directory with no unencrypted file costs 0 and never appears in the plan.
• Ties keep the finer plan (file calls), so the output does not touch
  more files than needed.

Part 3: no double encryption. If encrypting an already-encrypted file
corrupts it, a directory call is only legal when its subtree has zero
encrypted files. The recurrence keeps the same shape. Track the encrypted
count per subtree and allow option B only when that count is zero. With
file_cost = dir_cost = 1 this becomes the classic "smallest list of paths"
cover: name the highest directory that is entirely plaintext, otherwise list
the files.

from __future__ import annotations

from dataclasses import dataclass, field


@dataclass
class FileNode:
    name: str
    is_encrypted: bool


@dataclass
class DirectoryNode:
    name: str
    directories: list[DirectoryNode] = field(default_factory=list)
    files: list[FileNode] = field(default_factory=list)


@dataclass
class Plan:
    cost: int
    calls: list[str]          # paths passed to encrypt(), in DFS order
    plain: int                # unencrypted files in the subtree
    encrypted: int            # already-encrypted files in the subtree


def count_files(root: DirectoryNode) -> tuple[int, int]:
    """Part 1: (encrypted, unencrypted) across the whole tree, iteratively."""
    enc = plain = 0
    stack = [root]
    while stack:
        d = stack.pop()
        for f in d.files:
            if f.is_encrypted:
                enc += 1
            else:
                plain += 1
        stack.extend(d.directories)
    return enc, plain


def plan_encryption(
    root: DirectoryNode,
    file_cost: int,
    dir_cost: int,
    strict: bool = False,
) -> Plan:
    """Parts 2 and 3: cheapest set of encrypt() calls that leaves every file encrypted.

    encrypt(file) costs file_cost; encrypt(dir) costs dir_cost and encrypts every
    file in the subtree. With strict=True a directory call is only allowed when no
    file in its subtree is already encrypted (double encryption corrupts data).
    Post-order is done with an explicit stack so deep trees cannot overflow.
    """
    results: dict[int, Plan] = {}
    stack: list[tuple[DirectoryNode, str, bool]] = [(root, "/" + root.name, False)]
    while stack:
        node, path, children_done = stack.pop()
        if not children_done:
            stack.append((node, path, True))
            for child in node.directories:
                stack.append((child, f"{path}/{child.name}", False))
            continue

        # Option A: handle files here one by one and recurse into children.
        cost, calls = 0, []
        plain = encrypted = 0
        for f in node.files:
            if f.is_encrypted:
                encrypted += 1
            else:
                plain += 1
                cost += file_cost
                calls.append(f"{path}/{f.name}")
        for child in node.directories:
            sub = results.pop(id(child))
            cost += sub.cost
            calls.extend(sub.calls)
            plain += sub.plain
            encrypted += sub.encrypted

        # Option B: one call on this directory. Strictly cheaper only, so ties
        # keep the finer-grained plan; skip when nothing needs encrypting.
        allowed = plain > 0 and not (strict and encrypted > 0)
        if allowed and dir_cost < cost:
            cost, calls = dir_cost, [path]
        results[id(node)] = Plan(cost, calls, plain, encrypted)
    return results[id(root)]


if __name__ == "__main__":
    F = FileNode
    D = DirectoryNode
    tree = D("data", [
        D("logs", [], [F("a.log", False), F("b.log", False), F("c.log", False)]),
        D("keys", [], [F("k1", True), F("k2", False)]),
        D("tmp", [D("x", [], [F("y", True)])], []),
    ], [F("readme", False)])

    assert count_files(tree) == (2, 5)
    p = plan_encryption(tree, file_cost=1, dir_cost=2)
    assert p.cost == 2 and p.calls == ["/data"], p
    p = plan_encryption(tree, file_cost=1, dir_cost=2, strict=True)
    assert p.cost == 4 and p.calls == ["/data/readme", "/data/logs", "/data/keys/k2"], p
    p = plan_encryption(tree, file_cost=1, dir_cost=10)
    assert p.cost == 5 and len(p.calls) == 5, p
    p = plan_encryption(tree, file_cost=3, dir_cost=5)
    # logs: min(9,5)=5; keys: min(3,5)=3; tmp: 0; root: min(3+5+3, 5)=5
    assert p.cost == 5 and p.calls == ["/data"], p

    empty = D("e")
    assert count_files(empty) == (0, 0)
    assert plan_encryption(empty, 1, 1).calls == []
    done = D("r", [D("s", [], [F("f", True)])], [F("g", True)])
    assert plan_encryption(done, 1, 1).cost == 0

    # deep chain: no recursion limit
    deep = D("leaf", [], [F("f", False)])
    for i in range(5_000):
        deep = D(f"d{i}", [deep], [])
    p = plan_encryption(deep, file_cost=1, dir_cost=1)
    assert p.cost == 1 and p.calls[0].endswith("/leaf/f")
    print("ok")

Let N be the number of directories plus files.

• count_files: O(N) time, O(W) extra space for the stack (W is the
  widest frontier, at most N).
• plan_encryption: O(N) node visits. Each node does O(1) work apart
  from merging call lists and building paths. List merging is O(N) in
  total if you build the plan in a second pass from the stored decisions.
  Here it is O(N * depth) in the worst case, which is fine at interview
  size. Path strings add O(total path length).
• Space: O(N) for the per-node results (each is freed when its parent
  consumes it) and the explicit stack, so depth never hits a recursion limit.

Tree used below:

/data
  readme        (plain)
  logs/  a.log b.log c.log   (all plain)
  keys/  k1 (encrypted), k2 (plain)
  tmp/x/ y      (encrypted)

• count_files(tree) returns (2, 5).
• file_cost=1, dir_cost=2 returns cost 2, calls ["/data"]. One root call
  beats 5 file calls.
• file_cost=1, dir_cost=10 returns cost 5, with 5 file calls.
• file_cost=3, dir_cost=5 returns cost 5, calls ["/data"]. Here logs
  would cost 5, keys 3, tmp 0 and readme 3, so children(root) is 11.
• file_cost=1, dir_cost=2, strict=True returns cost 4, calls
  ["/data/readme", "/data/logs", "/data/keys/k2"]. You cannot call
  /data or /data/keys, because they contain encrypted files.
• An empty directory returns cost 0 and no calls.
• A tree where every file is already encrypted returns cost 0 and no calls,
  even when dir_cost < file_cost.
• A 5,000-deep chain with one plaintext leaf file returns cost 1 with no
  RecursionError.

• Per-file or size-based cost. encrypt(file) costs the file size, and
  encrypt(dir) costs a fixed overhead plus the total bytes under it,
  including files already encrypted. The recurrence is the same, with
  different leaf and option-B costs. Aggregate subtree byte totals in the same
  post-order pass.
• Mixed target states. Some files must stay plaintext (for example
  config or files with a do_not_encrypt flag). This works like strict mode:
  a directory call is legal only when its subtree contains no protected file.
  Count the protected files per subtree.
• Toggle semantics. Suppose encrypt(dir) flips every file in the
  subtree (an XOR cipher, or chmod-style toggles). Then the DP needs a
  parity state. best(d, p) is the cost when the ancestors have flipped an
  odd (p=1) or even number of times:
  best(d, p) = min(sum best(c, p), dir_cost + sum best(c, 1-p)).
  A file costs 0 if is_encrypted XOR p, otherwise file_cost.
• Incremental updates. Files are created and modified while the job runs.
  Cache best, plain and encrypted on each directory. A file change
  recomputes only the path to the root, which is O(depth) per update.
• Huge or remote trees. The tree comes from a paginated listing API. Stream
  it, keep only per-directory aggregates, and run the post-order step when the
  last child listing returns. Fan out listing calls with a bounded worker pool.
  Count calls, not CPU.

Mid-level (L4)

• Writes a correct traversal and counts for Part 1, and handles empty
  directories.
• Moves from "encrypt everything at the root" to seeing that the cost depends
  on the subtree. Gets the min(children, dir_cost) recurrence with some
  hints.
• Tests a small hand-built tree.

Senior (L5)

• States the recurrence unprompted and explains why subtrees are independent.
• Returns the actual plan, not just the cost, and handles zero-work subtrees
  and ties.
• Uses iterative post-order or names the recursion-depth risk. Adapts quickly
  to strict mode.

Staff (L6+)

• Treats the cost model as a parameter. Covers size-based costs, protected
  files and toggle parity with the same DP shape.
• Discusses incremental maintenance, remote or paginated listings, and
  failure mid-plan (idempotent calls, resuming from the plan).
• Avoids O(N * depth) list copying by storing per-node decisions and
  emitting the plan in a second top-down pass.