diff --git a/contest/src/main/java/com/github/contest/Execute.kt b/contest/src/main/java/com/github/contest/Execute.kt index a39e140..0ad3222 100644 --- a/contest/src/main/java/com/github/contest/Execute.kt +++ b/contest/src/main/java/com/github/contest/Execute.kt @@ -1,6 +1,8 @@ package com.github.contest +import com.github.contest.binaryTree.TreeNode +import com.github.contest.binaryTree.replaceValueInTree import com.github.contest.math.numberOfPowerfulInt import com.github.contest.slidingWindow.customStructure.rabinKarpMultiPattern import com.github.contest.slidingWindow.customStructure.slidingWindowClassic @@ -14,7 +16,15 @@ import java.util.TreeMap fun main() { - Example("Param") + // 5,4,9,1,10,null,7 + val root = TreeNode(5) + root?.left = TreeNode(4) + root?.right = TreeNode(9) + root.left?.left = TreeNode(1) + root?.left?.right = TreeNode(10) + root?.right?.right = TreeNode(7) + + replaceValueInTree(root).also { println(it) } } class Example(val param: String) { // Primary constructor parameter diff --git a/contest/src/main/java/com/github/contest/binaryTree/BinaryTreeLeetcode.kt b/contest/src/main/java/com/github/contest/binaryTree/BinaryTreeLeetcode.kt index 5301251..84afd39 100644 --- a/contest/src/main/java/com/github/contest/binaryTree/BinaryTreeLeetcode.kt +++ b/contest/src/main/java/com/github/contest/binaryTree/BinaryTreeLeetcode.kt @@ -156,3 +156,91 @@ fun evaluateTree(root: TreeNode?): Boolean { else -> left && right } } + +/** + * 2641. Cousins in Binary Tree II + */ + +fun replaceValueInTree(root: TreeNode?): TreeNode? { + root ?: return null + + val queue = LinkedList().apply { offer(root) } + val (nodeToParent, sumLevels, parentToChildrenSum) = bfsRelative(root) + var level = 0 + + while (queue.isNotEmpty()) { + val size = queue.size + val levelSum = sumLevels[level] + + for (i in 0 until size) { + val node = queue.poll() + + if (level == 0) { + node.`val` = 0 + } else { + val parent = nodeToParent[node]!! + val siblingsSum = parentToChildrenSum[parent]!! + node.`val` = levelSum - siblingsSum + } + + node?.left?.let { + queue.offer(it) + } + + node?.right?.let { + queue.offer(it) + } + } + + level++ + } + + return root +} + +private fun bfsRelative(root: TreeNode): TreeNodeRelative { + + val queue = LinkedList().apply { offer(root) } + val nodeToParent = mutableMapOf().apply { this[root] = null } + val parentToChildrenSum = mutableMapOf() + val sumLevels = mutableListOf() + + while (queue.isNotEmpty()) { + val size = queue.size + var levelSum = 0 + + for (i in 0 until size) { + val node = queue.poll() + var childsSum = 0 + levelSum += node.`val` + + node?.left?.let { + nodeToParent[it] = node + childsSum += it.`val` + queue.offer(it) + } + + node?.right?.let { + nodeToParent[it] = node + childsSum += it.`val` + queue.offer(it) + } + + parentToChildrenSum[node] = childsSum + } + + sumLevels.add(levelSum) + } + + return TreeNodeRelative( + nodeToParent = nodeToParent, + sumLevels = sumLevels, + parentToChildrenSum = parentToChildrenSum + ) +} + +private data class TreeNodeRelative( + val nodeToParent: Map, + val sumLevels: List, + val parentToChildrenSum: Map +)