Algorithm: topological sort collect coins in tree

DIRECTORY.md

@@ -145,6 +145,8 @@
       * [Test Cat And Mouse](https://github.com/BrianLusina/PythonSnips/blob/master/algorithms/graphs/cat_and_mouse/test_cat_and_mouse.py)
     * Cheapest Flights With K Stops
       * [Test Cheapest Flights With K Stops](https://github.com/BrianLusina/PythonSnips/blob/master/algorithms/graphs/cheapest_flights_with_k_stops/test_cheapest_flights_with_k_stops.py)
+    * Collect Coins In Tree
+      * [Test Collect Coins In Tree](https://github.com/BrianLusina/PythonSnips/blob/master/algorithms/graphs/collect_coins_in_tree/test_collect_coins_in_tree.py)
     * Course Schedule
       * [Test Course Schedule](https://github.com/BrianLusina/PythonSnips/blob/master/algorithms/graphs/course_schedule/test_course_schedule.py)
     * Evaluate Division

algorithms/graphs/collect_coins_in_tree/README.md

@@ -0,0 +1,61 @@
+# Collect Coins in a Tree
+
+There exists an undirected and unrooted tree with n nodes indexed from 0 to n - 1. You are given an integer n and a 2D
+integer array edges of length n - 1, where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi
+in the tree. You are also given an array coins of size n where coins[i] can be either 0 or 1, where 1 indicates the
+presence of a coin in the vertex i.
+
+Initially, you choose to start at any vertex in the tree. Then, you can perform the following operations any number of
+times: 
+
+- Collect all the coins that are at a distance of at most 2 from the current vertex, or 
+- Move to any adjacent vertex in the tree.
+
+Find the minimum number of edges you need to go through to collect all the coins and go back to the initial vertex.
+
+Note that if you pass an edge several times, you need to count it into the answer several times.
+
+## Examples
+
+Example 1:
+
+```text
+Input: coins = [1,0,0,0,0,1], edges = [[0,1],[1,2],[2,3],[3,4],[4,5]]
+Output: 2
+Explanation: Start at vertex 2, collect the coin at vertex 0, move to vertex 3, collect the coin at vertex 5 then move
+back to vertex 2.
+```
+
+Example 2:
+
+```text
+Input: coins = [0,0,0,1,1,0,0,1], edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[5,6],[5,7]]
+Output: 2
+Explanation: Start at vertex 0, collect the coins at vertices 4 and 3, move to vertex 2,  collect the coin at vertex 7,
+then move back to vertex 0.
+```
+
+## Constraints
+
+- n == coins.length
+- 1 <= n <= 3 * 10^4
+- 0 <= coins[i] <= 1
+- edges.length == n - 1
+- edges[i].length == 2
+- 0 <= ai, bi < n
+- ai != bi
+- edges represents a valid tree.
+
+## Hints
+
+- All leaves that do not have a coin are redundant and can be deleted from the tree.
+- Remove the leaves that do not have coins on them, so that the resulting tree will have a coin on every leaf.
+- In the remaining tree, remove each leaf node and its parent from the tree. The remaining nodes in the tree are the
+  ones that must be visited. Hence, the answer is equal to (# remaining nodes -1) * 2
+
+## Topics
+
+- Graph
+- Topological Sort
+- Trees
+- Arrays

algorithms/graphs/collect_coins_in_tree/__init__.py

@@ -0,0 +1,93 @@
+from typing import List
+from collections import defaultdict, deque
+
+
+def collect_the_coins(coins: List[int], edges: List[List[int]]) -> int:
+    # Build the adjacency list representation of the tree
+    graph = defaultdict(set)
+    for node_a, node_b in edges:
+        graph[node_a].add(node_b)
+        graph[node_b].add(node_a)
+
+    n = len(coins)
+
+    # Remove all leaf nodes that don't have coins. Initialize a queue with leaf nodes(degree 1) that have no coins
+    queue = deque(
+        node for node in range(n) if len(graph[node]) == 1 and coins[node] == 0
+    )
+
+    # Keep removing leaf nodes without coins
+    while queue:
+        current_node = queue.popleft()
+
+        # Remove this node from its neighbor's adjacency list
+        for neighbor in graph[current_node]:
+            graph[neighbor].remove(current_node)
+            # If neighbor becomes a leaf node and has no coin, add to queue
+            if coins[neighbor] == 0 and len(graph[neighbor]) == 1:
+                queue.append(neighbor)
+        # Clear the current node's connections
+        graph[current_node].clear()
+
+    # Remove two layers of lead nodes. This accounts for the collection distance of 2
+    for layer in range(2):
+        # Find all current leaf nodes
+        leaf_nodes = [node for node in range(n) if len(graph[node]) == 1]
+        # Remove all leaf nodes from the adjacency list
+        for leaf in leaf_nodes:
+            for neighbour in graph[leaf]:
+                graph[neighbour].remove(leaf)
+            graph[leaf].clear()
+
+    # Count remaining edgest that need to be traversed. An edge is counted if both its endpoints still exist in the graph
+    # Multiply by 2 because we need to traverse each edge twice (forward and back)
+    remaining_edges = sum(
+        len(graph[node_a]) > 0 and len(graph[node_b]) > 0 for node_a, node_b in edges
+    )
+
+    return remaining_edges * 2
+
+
+def collect_the_coins_2(coins: List[int], edges: List[List[int]]) -> int:
+    # Get the number of nodes
+    n = len(coins)
+
+    # Build adjacency set for each node (using sets for O(1) removal)
+    graph = [set() for _ in range(n)]
+    for a, b in edges:
+        graph[a].add(b)
+        graph[b].add(a)
+
+    # --- Phase 1: Topological sort to remove non-coin leaves ---
+    # Initialize queue with leaf nodes that have no coins
+    queue = deque()
+    for node in range(n):
+        # A leaf is a node with exactly one neighbor
+        if len(graph[node]) == 1 and coins[node] == 0:
+            queue.append(node)
+
+    # Repeatedly prune zero-coin leaves
+    while queue:
+        node = queue.popleft()
+        # For the single neighbor of this leaf
+        for neighbor in graph[node]:
+            graph[neighbor].discard(node)
+            # If neighbor becomes a no-coin leaf, add to queue
+            if len(graph[neighbor]) == 1 and coins[neighbor] == 0:
+                queue.append(neighbor)
+        # Remove all edges from this node
+        graph[node].clear()
+
+    # --- Phase 2: Prune two layers of leaves from the remaining tree ---
+    # First layer pruning: remove current coin-bearing leaves
+    for _ in range(2):
+        leaf_nodes = [node for node in range(n) if len(graph[node]) == 1]
+        for node in leaf_nodes:
+            for neighbor in graph[node]:
+                graph[neighbor].discard(node)
+            graph[node].clear()
+
+    # --- Result: count remaining edges * 2 ---
+    # Each remaining edge must be traversed twice (once each way)
+    remaining_edges = sum(len(neighbors) for neighbors in graph) // 2
+    return remaining_edges * 2

algorithms/graphs/collect_coins_in_tree/test_collect_coins_in_tree.py

@@ -0,0 +1,51 @@
+import unittest
+from typing import List
+
+from parameterized import parameterized
+from utils.test_utils import custom_test_name_func
+from algorithms.graphs.collect_coins_in_tree import (
+    collect_the_coins,
+    collect_the_coins_2,
+)
+
+COLLECT_COINS_IN_TREE_TEST_CASES = [
+    ([1, 0, 0, 0, 0, 1], [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5]], 2),
+    (
+        [0, 0, 0, 1, 1, 0, 0, 1],
+        [[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [5, 6], [5, 7]],
+        2,
+    ),
+    ([1], [], 0),
+    ([0, 0, 0, 0, 0], [[0, 1], [1, 2], [2, 3], [3, 4]], 0),
+    (
+        [1, 0, 0, 0, 0, 0, 0, 1],
+        [[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6], [6, 7]],
+        6,
+    ),
+    ([0, 1, 1, 1, 1], [[0, 1], [0, 2], [0, 3], [0, 4]], 0),
+    ([1, 1], [[0, 1]], 0),
+]
+
+
+class CollectCoinsInTreeTestCase(unittest.TestCase):
+    @parameterized.expand(
+        COLLECT_COINS_IN_TREE_TEST_CASES, name_func=custom_test_name_func
+    )
+    def test_collect_the_coins(
+        self, coins: List[int], edges: List[List[int]], expected: int
+    ):
+        actual = collect_the_coins(coins, edges)
+        self.assertEqual(expected, actual)
+
+    @parameterized.expand(
+        COLLECT_COINS_IN_TREE_TEST_CASES, name_func=custom_test_name_func
+    )
+    def test_collect_the_coins_2(
+        self, coins: List[int], edges: List[List[int]], expected: int
+    ):
+        actual = collect_the_coins_2(coins, edges)
+        self.assertEqual(expected, actual)
+
+
+if __name__ == "__main__":
+    unittest.main()