Thomas Ngo

Software Engineer

2 minute read

Problem Statement

leetcode problem link

Intuition

The problem involves navigating through a system of boxes that may be locked or unlocked, each potentially containing candies, keys to unlock other boxes, or additional boxes. The goal is to collect the maximum number of candies possible. This setup naturally leads to a simulation using a graph-like traversal while managing access via keys and status.

Approach

We simulate the process using a priority queue (heapq) to manage the order of boxes to process, prioritizing unlocked boxes. Here’s the step-by-step:

  1. Create a mapping from each box to its contained boxes (graph).
  2. Initialize a max-heap with all initialBoxes, storing them as (-status[box], box) so that unlocked boxes (status = 1) are prioritized.
  3. Maintain a visited set to ensure each box is processed only once.
  4. While the heap is not empty:
    • Pop the box with the highest priority.
    • If it is accessible (status[box] == 1) and not yet visited:
      • Mark it as visited.
      • Add its candies to the result.
      • For each key in this box, update the status of the corresponding box to unlocked.
      • For each contained box, push it into the heap with its current status.

This ensures all accessible paths are explored, keys are propagated properly, and boxes are not reprocessed.

Complexity

  • Time complexity:
    \(O(n \cdot \log n + k)\)
    Where (n) is the number of boxes and (k) is the total number of keys and contained boxes. Each box is pushed and popped at most once due to the visited check.

  • Space complexity:
    \(O(n + k)\)
    For the visited set, heap, and auxiliary structures like the graph and key/box lists.

Code

import heapq

class Solution:
    def maxCandies(self, status: List[int], candies: List[int], keys: List[List[int]], containedBoxes: List[List[int]], initialBoxes: List[int]) -> int:
        graph = {i: box for i, box in enumerate(containedBoxes)}
        max_heap = []
        for initialBox in initialBoxes:
            heapq.heappush(max_heap, [-1 * status[initialBox], initialBox])
        res = 0
        visited = set()
        while max_heap:
            s, box_label = heapq.heappop(max_heap)
            s *= -1
            if status[box_label] == 1 and box_label not in visited:
                visited.add(box_label)
                res += candies[box_label]
                for key in keys[box_label]:
                    status[key] = 1
                for nei in graph[box_label]:
                    heapq.heappush(max_heap, [-1 * status[nei], nei])
        return res

Editorial

class Solution:
    def maxCandies(
        self,
        status: List[int],
        candies: List[int],
        keys: List[List[int]],
        containedBoxes: List[List[int]],
        initialBoxes: List[int],
    ) -> int:
        n = len(status)
        can_open = [status[i] == 1 for i in range(n)]
        has_box, used = [False] * n, [False] * n

        q = collections.deque()
        ans = 0
        for box in initialBoxes:
            has_box[box] = True
            if can_open[box]:
                q.append(box)
                used[box] = True
                ans += candies[box]

        while len(q) > 0:
            big_box = q.popleft()
            for key in keys[big_box]:
                can_open[key] = True
                if not used[key] and has_box[key]:
                    q.append(key)
                    used[key] = True
                    ans += candies[key]
            for box in containedBoxes[big_box]:
                has_box[box] = True
                if not used[box] and can_open[box]:
                    q.append(box)
                    used[box] = True
                    ans += candies[box]

        return ans

You may also enjoy