From c5081ca25fcf38dab5becae38f3e5dfe3de3dc30 Mon Sep 17 00:00:00 2001 From: BrianLusina <12752833+BrianLusina@users.noreply.github.com> Date: Thu, 20 Aug 2026 10:06:02 +0300 Subject: [PATCH 1/2] feat(algorithms, graphs, topological-sort): collect coins in a tree: Collect coins in a tree using topological sort algorithm pattern --- .../graphs/collect_coins_in_tree/README.md | 61 ++++++++++++ .../graphs/collect_coins_in_tree/__init__.py | 93 +++++++++++++++++++ .../test_collect_coins_in_tree.py | 51 ++++++++++ 3 files changed, 205 insertions(+) create mode 100644 algorithms/graphs/collect_coins_in_tree/README.md create mode 100644 algorithms/graphs/collect_coins_in_tree/__init__.py create mode 100644 algorithms/graphs/collect_coins_in_tree/test_collect_coins_in_tree.py diff --git a/algorithms/graphs/collect_coins_in_tree/README.md b/algorithms/graphs/collect_coins_in_tree/README.md new file mode 100644 index 00000000..1fcc482b --- /dev/null +++ b/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 diff --git a/algorithms/graphs/collect_coins_in_tree/__init__.py b/algorithms/graphs/collect_coins_in_tree/__init__.py new file mode 100644 index 00000000..1f9cfbc2 --- /dev/null +++ b/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 diff --git a/algorithms/graphs/collect_coins_in_tree/test_collect_coins_in_tree.py b/algorithms/graphs/collect_coins_in_tree/test_collect_coins_in_tree.py new file mode 100644 index 00000000..17a3664f --- /dev/null +++ b/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() From be67052d87eb3b9a437fc98bdf3975bb956d0731 Mon Sep 17 00:00:00 2001 From: github-actions <${GITHUB_ACTOR}@users.noreply.github.com> Date: Thu, 20 Aug 2026 07:07:02 +0000 Subject: [PATCH 2/2] updating DIRECTORY.md --- DIRECTORY.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/DIRECTORY.md b/DIRECTORY.md index 84cf928e..38439fe6 100644 --- a/DIRECTORY.md +++ b/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