How to “draw” a Binary Tree in Console? - swift

Question

How can I print a binary tree in Swift so that the input 79561 prints output like this:
7
/ \
5 9
/ \
1 6
I tried to arrange this with some code using For Loops and If Statements but it didn't worked.
My code is:
import UIKit
//Variable "node" used only to arrange it in output.
var node = "0"
var space = " "
var linkLeft = "/"
var linkRight = "\\"
var str = "Hello, playground"
var height = 6
var width = height * 2 + 1
print()
//Height
for h in 1...height {
//Width
for w in 1...width {
switch h {
case 1:
if(w == width/2 + h) {
print(node, terminator: "")
} else {
print(space, terminator: "")
}
if (w == width) {
print()
}
case 2:
//print(linkLeft, terminator: "")
if(w == width/3 + h) {
print(linkLeft, terminator: "")
} else if(w == width/3 + h + 4) {
print(linkRight, terminator: "")
} else {
print(space, terminator: "")
}
if (w == width) {
print()
}
case 3:
if(w == width/5 + h) {
print(node, terminator: "")
} else if(w == width/h + h){
print(node, terminator: "")
} else {
print(space, terminator: "")
}
if (w == width) {
print()
}
break
default:
break
}
}
}
I tried to use two For Loops one for height and other one for width. But it's not working if number of nodes changes. For now I just tried to arrange places of links (/ and \), nodes and spaces, so it's not working.
Is there a possible way to do this ?

Answer 1

First you have to define a hierarchical tree structure (class) that allows recursive traversal of the tree nodes. How you implement it doesn't matter as long as it can provide a descriptive string and access to its left and right sub nodes.
For example (I used this for testing purposes):
class TreeNode
{
var value : Int
var left : TreeNode? = nil
var right : TreeNode? = nil
init(_ rootValue:Int)
{ value = rootValue }
#discardableResult
func addValue( _ newValue:Int) -> TreeNode
{
if newValue == value // exclude duplicate entries
{ return self }
else if newValue < value
{
if let newNode = left?.addValue(newValue)
{ return newNode }
left = TreeNode(newValue)
return left!
}
else
{
if let newNode = right?.addValue(newValue)
{ return newNode }
right = TreeNode(newValue)
return right!
}
}
}
Then you can create a recursive function to obtain the lines to print. Each line will need to be aware of lower level lines so the list of lines needs to be built from the bottom up. Recursion is an easy way to achieve this kind of interdependency.
Here's an example of a generic function that will work for any binary tree class. It expects a root node and a function (or closure) to access a node's description and left/right sub nodes :
public func treeString<T>(_ node:T, reversed:Bool=false, isTop:Bool=true, using nodeInfo:(T)->(String,T?,T?)) -> String
{
// node value string and sub nodes
let (stringValue, leftNode, rightNode) = nodeInfo(node)
let stringValueWidth = stringValue.count
// recurse to sub nodes to obtain line blocks on left and right
let leftTextBlock = leftNode == nil ? []
: treeString(leftNode!,reversed:reversed,isTop:false,using:nodeInfo)
.components(separatedBy:"\n")
let rightTextBlock = rightNode == nil ? []
: treeString(rightNode!,reversed:reversed,isTop:false,using:nodeInfo)
.components(separatedBy:"\n")
// count common and maximum number of sub node lines
let commonLines = min(leftTextBlock.count,rightTextBlock.count)
let subLevelLines = max(rightTextBlock.count,leftTextBlock.count)
// extend lines on shallower side to get same number of lines on both sides
let leftSubLines = leftTextBlock
+ Array(repeating:"", count: subLevelLines-leftTextBlock.count)
let rightSubLines = rightTextBlock
+ Array(repeating:"", count: subLevelLines-rightTextBlock.count)
// compute location of value or link bar for all left and right sub nodes
// * left node's value ends at line's width
// * right node's value starts after initial spaces
let leftLineWidths = leftSubLines.map{$0.count}
let rightLineIndents = rightSubLines.map{$0.prefix{$0==" "}.count }
// top line value locations, will be used to determine position of current node & link bars
let firstLeftWidth = leftLineWidths.first ?? 0
let firstRightIndent = rightLineIndents.first ?? 0
// width of sub node link under node value (i.e. with slashes if any)
// aims to center link bars under the value if value is wide enough
//
// ValueLine: v vv vvvvvv vvvvv
// LinkLine: / \ / \ / \ / \
//
let linkSpacing = min(stringValueWidth, 2 - stringValueWidth % 2)
let leftLinkBar = leftNode == nil ? 0 : 1
let rightLinkBar = rightNode == nil ? 0 : 1
let minLinkWidth = leftLinkBar + linkSpacing + rightLinkBar
let valueOffset = (stringValueWidth - linkSpacing) / 2
// find optimal position for right side top node
// * must allow room for link bars above and between left and right top nodes
// * must not overlap lower level nodes on any given line (allow gap of minSpacing)
// * can be offset to the left if lower subNodes of right node
// have no overlap with subNodes of left node
let minSpacing = 2
let rightNodePosition = zip(leftLineWidths,rightLineIndents[0..<commonLines])
.reduce(firstLeftWidth + minLinkWidth)
{ max($0, $1.0 + minSpacing + firstRightIndent - $1.1) }
// extend basic link bars (slashes) with underlines to reach left and right
// top nodes.
//
// vvvvv
// __/ \__
// L R
//
let linkExtraWidth = max(0, rightNodePosition - firstLeftWidth - minLinkWidth )
let rightLinkExtra = linkExtraWidth / 2
let leftLinkExtra = linkExtraWidth - rightLinkExtra
// build value line taking into account left indent and link bar extension (on left side)
let valueIndent = max(0, firstLeftWidth + leftLinkExtra + leftLinkBar - valueOffset)
let valueLine = String(repeating:" ", count:max(0,valueIndent))
+ stringValue
let slash = reversed ? "\\" : "/"
let backSlash = reversed ? "/" : "\\"
let uLine = reversed ? "¯" : "_"
// build left side of link line
let leftLink = leftNode == nil ? ""
: String(repeating: " ", count:firstLeftWidth)
+ String(repeating: uLine, count:leftLinkExtra)
+ slash
// build right side of link line (includes blank spaces under top node value)
let rightLinkOffset = linkSpacing + valueOffset * (1 - leftLinkBar)
let rightLink = rightNode == nil ? ""
: String(repeating: " ", count:rightLinkOffset)
+ backSlash
+ String(repeating: uLine, count:rightLinkExtra)
// full link line (will be empty if there are no sub nodes)
let linkLine = leftLink + rightLink
// will need to offset left side lines if right side sub nodes extend beyond left margin
// can happen if left subtree is shorter (in height) than right side subtree
let leftIndentWidth = max(0,firstRightIndent - rightNodePosition)
let leftIndent = String(repeating:" ", count:leftIndentWidth)
let indentedLeftLines = leftSubLines.map{ $0.isEmpty ? $0 : (leftIndent + $0) }
// compute distance between left and right sublines based on their value position
// can be negative if leading spaces need to be removed from right side
let mergeOffsets = indentedLeftLines
.map{$0.count}
.map{leftIndentWidth + rightNodePosition - firstRightIndent - $0 }
.enumerated()
.map{ rightSubLines[$0].isEmpty ? 0 : $1 }
// combine left and right lines using computed offsets
// * indented left sub lines
// * spaces between left and right lines
// * right sub line with extra leading blanks removed.
let mergedSubLines = zip(mergeOffsets.enumerated(),indentedLeftLines)
.map{ ( $0.0, $0.1, $1 + String(repeating:" ", count:max(0,$0.1)) ) }
.map{ $2 + String(rightSubLines[$0].dropFirst(max(0,-$1))) }
// Assemble final result combining
// * node value string
// * link line (if any)
// * merged lines from left and right sub trees (if any)
let treeLines = [leftIndent + valueLine]
+ (linkLine.isEmpty ? [] : [leftIndent + linkLine])
+ mergedSubLines
return (reversed && isTop ? treeLines.reversed(): treeLines)
.joined(separator:"\n")
}
To actually print, you'll need to supply the function with your class's node and a closure to access node descriptions and the left and right sub nodes.
extension TreeNode
{
var asString:String { return treeString(self){("\($0.value)",$0.left,$0.right)} }
}
var root = TreeNode(7)
root.addValue(9)
root.addValue(5)
root.addValue(6)
root.addValue(1)
print(root.asString)
// 7
// / \
// 5 9
// / \
// 1 6
//
root = TreeNode(80)
root.addValue(50)
root.addValue(90)
root.addValue(10)
root.addValue(60)
root.addValue(30)
root.addValue(70)
root.addValue(55)
root.addValue(5)
root.addValue(35)
root.addValue(85)
print(root.asString)
// 80
// ___/ \___
// 50 90
// __/ \__ /
// 10 60 85
// / \ / \
// 5 30 55 70
// \
// 35
//
[EDIT] Improved logic to use less space on trees with deeper right side than left. Cleaned up code and added comments to explain how it works.
//
// 12
// / \
// 10 50
// / __/ \__
// 5 30 90
// \ /
// 35 70
// / \
// 60 85
// /
// 55
//
// 12
// / \
// 10 30
// / \
// 5 90
// /
// 85
// /
// 70
// /
// 55
// /
// 48
// /
// 45
// /
// 40
// /
// 35
//
[EDIT2] made the function generic and adapted explanations.
With the generic function, the data doesn't even need to be in an actual tree structure.
For example, you could print an array containing a heap tree:
extension Array
{
func printHeapTree(reversed:Bool = false)
{
let tree = treeString( 0, reversed:reversed )
{
let left = { $0 < self.count ? $0 : nil}($0 * 2 + 1)
let right = { $0 < self.count ? $0 : nil}($0 * 2 + 2)
return ( "\(self[$0])", left, right )
}
print(tree)
}
}
let values = [7,5,9,1,6]
values.printHeapTree()
// 7
// / \
// 5 9
// / \
// 1 6
let family = [ "Me","Paul","Rosa","Vincent","Jody","John","Kate"]
family.printHeapTree()
// Me
// ___/ \___
// Paul Rosa
// / \ / \
// Vincent Jody John Kate
But for that last example, a family tree is usually presented upside down. So, I adjusted the function to allow printing a reversed tree:
family.printHeapTree(reversed:true)
// Vincent Jody John Kate
// \ / \ /
// Paul Rosa
// ¯¯¯\ /¯¯¯
// Me
[EDIT3] Added condition to exclude duplicate entries from tree in example class (TreeNode) per Emm's request
[EDIT4] changed formula for mergedSubLines so that it will compile in an actual project (was testing this in the playground).
[EDIT5] minor adjustments for Swift4, added ability to print a reversed tree, changed Array example to a heap tree.

Answer 2

My easy way that I found and modified a bit for swift
import UIKit
// Binary tree by class
class TreeNode
{
var value : Int
var left : TreeNode? = nil
var right : TreeNode? = nil
init(_ rootValue:Int) {
value = rootValue
}
#discardableResult
func addValue( _ newValue:Int) -> TreeNode
{
if newValue == value {
return self
}
else if newValue < value {
if let newNode = left?.addValue(newValue) {
return newNode
}
left = TreeNode(newValue)
return left!
}
else {
if let newNode = right?.addValue(newValue) {
return newNode
}
right = TreeNode(newValue)
return right!
}
}
func myPrint () {
printTree(prefix: "", n: self, isLeft: false)
}
}
func printTree(prefix: String, n: TreeNode, isLeft: Bool) {
print(prefix, (isLeft ? "|-- " : "\\-- "), n.value)
if n.left != nil {
printTree(prefix: "\(prefix) \(isLeft ? "| " : " ") ", n: n.left!, isLeft: true)
}
if n.right != nil {
printTree(prefix: "\(prefix) \(isLeft ? "| " : " ") ", n: n.right!, isLeft: false)
}
}
Input
var root = TreeNode(80)
root.addValue(50)
root.addValue(90)
root.addValue(10)
root.addValue(60)
root.addValue(30)
root.addValue(70)
root.addValue(55)
root.addValue(5)
root.addValue(35)
root.addValue(85)
root.addValue(84)
root.addValue(86)
root.addValue(92)
root.addValue(100)
root.myPrint()
Output
\-- 80
|-- 50
| |-- 10
| | |-- 5
| | \-- 30
| | \-- 35
| \-- 60
| |-- 55
| \-- 70
\-- 90
|-- 85
| |-- 84
| \-- 86
\-- 92
\-- 100