diff --git a/DIRECTORY.md b/DIRECTORY.md index 5f314c317..74243cd06 100644 --- a/DIRECTORY.md +++ b/DIRECTORY.md @@ -173,6 +173,7 @@ * [Binary Tree Path Sum](data_structures/binary_tree/binary_tree_path_sum.py) * [Binary Tree Traversals](data_structures/binary_tree/binary_tree_traversals.py) * [Diff Views Of Binary Tree](data_structures/binary_tree/diff_views_of_binary_tree.py) + * [Distribute Coins](data_structures/binary_tree/distribute_coins.py) * [Fenwick Tree](data_structures/binary_tree/fenwick_tree.py) * [Inorder Tree Traversal 2022](data_structures/binary_tree/inorder_tree_traversal_2022.py) * [Is Bst](data_structures/binary_tree/is_bst.py) diff --git a/data_structures/binary_tree/distribute_coins.py b/data_structures/binary_tree/distribute_coins.py new file mode 100644 index 000000000..ea02afc2c --- /dev/null +++ b/data_structures/binary_tree/distribute_coins.py @@ -0,0 +1,135 @@ +""" +Author : Alexander Pantyukhin +Date : November 7, 2022 + +Task: +You are given a tree root of a binary tree with n nodes, where each node has +node.data coins. There are exactly n coins in whole tree. + +In one move, we may choose two adjacent nodes and move one coin from one node +to another. A move may be from parent to child, or from child to parent. + +Return the minimum number of moves required to make every node have exactly one coin. + +Example 1: + + 3 + / \ + 0 0 + +Result: 2 + +Example 2: + + 0 + / \ + 3 0 + +Result 3 + +leetcode: https://leetcode.com/problems/distribute-coins-in-binary-tree/ + +Implementation notes: +User depth-first search approach. + +Let n is the number of nodes in tree +Runtime: O(n) +Space: O(1) +""" + +from __future__ import annotations + +from collections import namedtuple +from dataclasses import dataclass + + +@dataclass +class TreeNode: + data: int + left: TreeNode | None = None + right: TreeNode | None = None + + +CoinsDistribResult = namedtuple("CoinsDistribResult", "moves excess") + + +def distribute_coins(root: TreeNode | None) -> int: + """ + >>> distribute_coins(TreeNode(3, TreeNode(0), TreeNode(0))) + 2 + >>> distribute_coins(TreeNode(0, TreeNode(3), TreeNode(0))) + 3 + >>> distribute_coins(TreeNode(0, TreeNode(0), TreeNode(3))) + 3 + >>> distribute_coins(None) + 0 + >>> distribute_coins(TreeNode(0, TreeNode(0), TreeNode(0))) + Traceback (most recent call last): + ... + ValueError: The nodes number should be same as the number of coins + >>> distribute_coins(TreeNode(0, TreeNode(1), TreeNode(1))) + Traceback (most recent call last): + ... + ValueError: The nodes number should be same as the number of coins + """ + + if root is None: + return 0 + + # Validation + def count_nodes(node: TreeNode | None) -> int: + """ + >>> count_nodes(None): + 0 + """ + if node is None: + return 0 + + return count_nodes(node.left) + count_nodes(node.right) + 1 + + def count_coins(node: TreeNode | None) -> int: + """ + >>> count_coins(None): + 0 + """ + if node is None: + return 0 + + return count_coins(node.left) + count_coins(node.right) + node.data + + if count_nodes(root) != count_coins(root): + raise ValueError("The nodes number should be same as the number of coins") + + # Main calculation + def get_distrib(node: TreeNode | None) -> CoinsDistribResult: + """ + >>> get_distrib(None) + namedtuple("CoinsDistribResult", "0 2") + """ + + if node is None: + return CoinsDistribResult(0, 1) + + left_distrib_moves, left_distrib_excess = get_distrib(node.left) + right_distrib_moves, right_distrib_excess = get_distrib(node.right) + + coins_to_left = 1 - left_distrib_excess + coins_to_right = 1 - right_distrib_excess + + result_moves = ( + left_distrib_moves + + right_distrib_moves + + abs(coins_to_left) + + abs(coins_to_right) + ) + result_excess = node.data - coins_to_left - coins_to_right + + return CoinsDistribResult(result_moves, result_excess) + + return get_distrib(root)[0] + + +if __name__ == "__main__": + import doctest + + doctest.testmod()