Sometimes, our trees can grow a little out of hand with branches and leaves that we don’t want anymore. Now that we’ve defined equality in our TreeNodes, we can start the process of pruning out any unneeded information. Our pruning process will be recursive in order to work our way down every branch and leaf of the structure, and it will also take advantage of some of the built-in Array methods already included in the Swift library. As with all recursion, the first step is to identify our base case to return out of. In our case the base cases will be:
- The
TreeNode
has no children, so return out of the method - The
TreeNode
’s children contain the item to remove, so remove the child and return
If neither of those cases is true, we call the removeChild(_:)
method on each child
in the list of children
.
Instructions
Below the addChild(_:)
methods, create a new function, removeChild(_:)
that has one parameter, a TreeNode
called nodeToRemove
, omit the argument label and leave the function body blank.
Inside the method create an if
statement that satisfies our first base case, checking if the node’s children
Array is empty.
Remember this is checking the node that the method was called on, not the node passed as an argument, therefore we can reference it directly as an instance property.
If it isEmpty
, return
out of the method.
Add an else if
statement that checks if children
.contains()
the node to remove. Inside the else if
, use the .removeAll()
method of children
. This method takes a closure, named where
, as a parameter and iterates over each item in the list, executing the where
closure. For each iteration where
the closure returns true
, it removes the item.
We will be using trailing closure syntax, a technique that can be used when the last argument of a function is a closure. For now, there should be no arguments in the method call and a set of open and closed curly braces directly after the method, .someMethod() {}
.
Let’s work on that closure!
The closure should return whether the element being iterated on is equal to nodeToRemove
. Because TreeNode
conforms to Equatable
, we can use the ==
operator. Use the $0
shorthand notation.
After the .removeAll()
call and still inside the else if
statement, return
out of the method.
Add a final else
statement to start our recursive logic. Inside, loop through each child
in children
and call the removeChild()
method on each child, passing in nodeToRemove
as the argument for the recursive removeChild()
method.
Test your code! It looks like someone accidentally assigned Fido as one of Clifford’s children, he should only have one child, Max.
Remove Fido (puppy1
) from Clifford on the second to last line of code in the editor using your brand new removeChild(_:)
function and then uncomment the print statement on the last line. You should see Clifford’s children decrease by one.