Working with binary trees and drawing them (p5js canvas library)

generateAndDrawTree1.js

/////////////////////////////////////////////////////////////////////////////////////
// GENERATING AND DRAWING TREES OF n LEVELS -- for now, easy mode -- perfect binary trees
// Note: using p5js for drawing functions; didn't include p5js setup() here
/////////////////////////////////////////////////////////////////////////////////////


///////////////////////////////////////////////
// SETUP / INITIALIZING VARIBALES
///////////////////////////////////////////////

// Iterate step by step with mouse clicks
document.addEventListener("click", generateAndDrawTree);

// PARAMETER -- tree with n levels:
let numLevels = 4;

// Setup for tree traversal:
let totalNodes = 2**numLevels - 1;     // NOTE: assuming the tree will always have a root node
let tree = {};
let node = tree;
let count = 1;
let nextNodes = [];

// Setup for drawing:
const canvasWidth = 900, canvasHeight = 900;   // canvas dimensions
const nodeSize = 50;                            // given as diameter of each node
const xMargin = 10;                         // minimum X margin between nodes
const yMargin = 25;                         // Y margin between each level

// Initial distance between nodes is the minimum distance doubled for each level (aside from bottom level and root level)
let nodeDistance = 2**(numLevels - 1) * (nodeSize + xMargin);

let startOfNextLevel = 1; // functions as the "outer index" (replacing nested loop from previous version)
let x1 = canvasWidth/2;
let currentX;
let y = nodeSize/2 + yMargin;

// For drawing each level a different color:
let colorArray = ["#ffffba", "#baffc9", "#bae1ff", "#ffdfba", "#ffb3ba"];

///////////////////////////////////////////////////////////////////////////////////////////////////////////////

function generateAndDrawTree() {

  console.log("count: " + count);
  console.log("startOfNextLevel: " + startOfNextLevel);
  console.log("x: " + currentX);
  console.log("y: " + y);
    
  // For the root node only (this is to visualize the root node as its own independent step)
  if (count === 1) {
    node.value = "01";
    
    // DRAWING:
    currentX = x1;
    fill( colorArray[0] ); // using log2(startOfNextLevel) would also work as an index here. yay logarithms! :)
    ellipse(currentX, y, nodeSize, nodeSize);
    node.x = currentX;
    node.y = y;
    
    // Iterate
    count++;
    startOfNextLevel *= 2;

    console.log( JSON.stringify(tree, null, '\t' ) ); 

    return;
  }
  
  // "outer loop" -- for each new level of nodes:
  if (count === startOfNextLevel) { 

    fill( colorArray[Math.log2(startOfNextLevel)] );
    
    // For each subsequent level, distance between nodes is halved as we move DOWN a level (where top level is level 1)
    nodeDistance /= 2;

    // Update the y coordinate for the next level down
    y += nodeSize + yMargin;

    // Update x coordinate of left-most node for the next level down:
    x1 -= nodeDistance/2;
    
    currentX = x1;
    
    startOfNextLevel *= 2;   // the count/value of first node of each level doubles: 1, 2, 4, 8, 16 ...

  }
  
  
  // "inner loop" -- generate and draw each node
  if (count <= totalNodes) {
    
    // Conditions only needed to add one node per iteration:
    if (!node.leftChild) {
      
      node.leftChild = {value: count.toString(2)};
      nextNodes.push(node.leftChild);
      count++;
            
      // DRAWING:
      ellipse(currentX, y, nodeSize, nodeSize);
      node.leftChild.x = currentX;
      node.leftChild.y = y;
      currentX = currentX + nodeDistance;
      
    } else if (!node.rightChild) {
      
      node.rightChild = {value: count.toString(2)};
      nextNodes.push(node.rightChild);
      count++;
      
      // DRAWING:
      ellipse(currentX, y, nodeSize, nodeSize);
      node.rightChild.x = currentX;
      node.rightChild.y = y;
      currentX = currentX + nodeDistance;

      node = nextNodes.shift()
    }

    console.log( JSON.stringify(tree, null, '\t' ) ); 

  } // end looping condition
	
} // end of function

generateRandomTreeWithCoords.js

////////////////////////////////////////////////////////////////////////////
// Generate a NON-PERFECT binary tree with randomly placed node
// of a given depth and total number of nodes
////////////////////////////////////////////////////////////////////////////

// NOTE: this version runs all at once, not step-by-step like the other version of the tree generating function

function generateRandomTreeWithCoords(numLevels) {

  // Initialize with the tree itself as the first node
  let tree = {};
  let node = tree;
  let nextNodes = [];

  // Initial distance between nodes is the minimum distance doubled for each level (aside from bottom level and root level)
  let nodeDistance = 2**(numLevels - 1) * (nodeSize + xMargin);  // this time: initialize before the loop, half it at end of outer loops
  let level = 1;
  let count = 1;
  let numNodes = 1;
  
  let initialX = canvasWidth/2;
  let y = nodeSize/2 + yMargin;
  
  // SET UP ROOT NODE:
  node.value = numNodes;
  node.x = initialX;
  node.y = y;
  node.level = level;
  count++;
  numNodes++;
  
  // GO TO NEXT LEVEL:
  level++;
  nodeDistance /= 2; // For each subsequent level, distance between nodes is halved as we move DOWN a level (where top level is level 1)
  y += nodeSize + yMargin; // Update the y coordinate for the next level down

  console.log( JSON.stringify(tree, null, '\t' ) ); 
  
  while (node && level <= numLevels) {
    
    console.log("level: " + level);
    console.log("count: " + count);
    console.log("THIS NODE: " + node.value + " -- " + node.x + ", " + node.y);
    
    if ( Math.round(Math.random()) ) {

      node.leftChild = {value: numNodes};
      node.leftChild.x = node.x - nodeDistance/2;
      node.leftChild.y = y;
      node.leftChild.level = level;
      nextNodes.push(node.leftChild);
      numNodes++;
    }
    count++;
    
    if ( Math.round(Math.random()) ) {
      
      node.rightChild = {value: numNodes};
      node.rightChild.x = node.x + nodeDistance/2;
      node.rightChild.y = y;
      node.rightChild.level = level;
      nextNodes.push(node.rightChild);
      numNodes++;
    }
    count++;

    // GO TO NEXT LEVEL:
    node = nextNodes.shift(); // iterate through nodes
    level++;
    nodeDistance /= 2; // For each subsequent level, distance between nodes is halved as we move DOWN a level (where top level is level 1)
    y += nodeSize + yMargin; // Update the y coordinate for the next level down

    console.log( JSON.stringify(tree, null, '\t' ) ); 

  } // end looping condition
	
  return tree;
  
} // end of function

generateTree1.js

/////////////////////////////////////////////////////////////////////////////////////
// GENERATING TREES OF n LEVELS -- for now, easy mode -- perfect binary trees
/////////////////////////////////////////////////////////////////////////////////////


//////////////////////////////////////////
// SETUP / INITIALIZING VARIBALES
/////////////////////////////////////////////////////////////////////////////////////

// Iterate step by step with mouse clicks
// ***** Use any of the functions defined below as the callback here: ********
document.addEventListener("click", generateTree);

// NOTE: To run the code below as a normal loop (instead of stepping through),
// just replace the word "if" with "while" in the code, and you're good to go!

// PARAMETER -- tree with n levels:
let numLevels = 4;

// Initialize with the tree itself as the first node
let totalNodes = 2**numLevels - 1;     // NOTE: assuming the tree will always have a root node
let tree = {};
let node = tree;
let count = 1;
let nextNodes = [];

////////////////////////////////////////////////////////////////////////////

function generateTree() {

  // For iterating step by step -- this way, creating the root node is step 1!
  if (count === 1) {
    node.value = "01";
    count++;

    console.log( JSON.stringify(tree, null, '\t' ) ); 

    return;
  }
  
  if (count <= totalNodes) {

    // Conditions only needed to add one node per iteration:
    if (!node.leftChild) {
      node.leftChild = {value: count.toString(2)};
      nextNodes.push(node.leftChild);
      count++;
    } else if (!node.rightChild) {
      node.rightChild = {value: count.toString(2)};
      nextNodes.push(node.rightChild);
      count++;

      node = nextNodes.shift()
    }

    console.log( JSON.stringify(tree, null, '\t' ) ); 

  } // end looping condition
	
} // end of function

generateTreeWithCoords.js

// NOTE: this version runs all at once, not step-by-step like the other version of the tree generating function

function generateTreeWithCoords(numLevels) {

  
  // Initialize with the tree itself as the first node
  let totalNodes = 2**numLevels - 1;     // NOTE: assuming the tree will always have a root node
  let tree = {};
  let node = tree;
  let count = 1;
  let nextNodes = [];


  // Initial distance between nodes is the minimum distance doubled for each level (aside from bottom level and root level)
  let nodeDistance = 2**(numLevels - 1) * (nodeSize + xMargin);  // this time: initialize before the loop, half it at end of outer loops
  let level = 1;
  let startOfNextLevel = 1; // functions as the "outer index" (replacing nested loop from previous version)
  
  
  let x1 = canvasWidth/2;
  let currentX;
  let y = nodeSize/2 + yMargin;

  // For the root node only:
  currentX = x1;
  
  node.value = "01";
  node.x = currentX;
  node.y = y;
  node.level = level;
  
  // Iterate
  count++;
  startOfNextLevel *= 2;

  console.log( JSON.stringify(tree, null, '\t' ) ); 
  
  
  while (count <= totalNodes) {
    
    console.log("count: " + count);
    console.log("startOfNextLevel: " + startOfNextLevel);
    console.log("x: " + currentX);
    console.log("y: " + y);
    
    // "outer loop" -- for each new level of nodes:
    if (count === startOfNextLevel) { 

      // For each subsequent level, distance between nodes is halved as we move DOWN a level (where top level is level 1)
      nodeDistance /= 2;

      // Update the y coordinate for the next level down
      y += nodeSize + yMargin;

      // Update x coordinate of left-most node for the next level down:
      x1 -= nodeDistance/2;

      currentX = x1;
      level++;
      startOfNextLevel *= 2;   // the count/value of first node of each level doubles: 1, 2, 4, 8, 16 ...

    }
    
    node.leftChild = {value: count.toString(2)};
    node.leftChild.x = currentX;
    node.leftChild.y = y;
    node.leftChild.level = level;
    nextNodes.push(node.leftChild);
    count++;
    
    currentX = currentX + nodeDistance;

    node.rightChild = {value: count.toString(2)};
    node.rightChild.x = currentX;
    node.rightChild.y = y;
    node.rightChild.level = level;
    nextNodes.push(node.rightChild);
    count++;

    currentX = currentX + nodeDistance;
    node = nextNodes.shift();


    console.log( JSON.stringify(tree, null, '\t' ) ); 

  } // end looping condition
	
  return tree;
  
} // end of function

getTreeHeight.js

/////////////////////////////////////////////////////////////////////////////////////
// TRAVERSING ANY BINARY TREE TO FIND ITS HEIGHT / NUM OF LEVELS
// First as a step-through function, then as a regular function with output
/////////////////////////////////////////////////////////////////////////////////////


/////////////////////////////////////////////////////////////////////////////////////
// EXAMPLE TREES -- a couple perfect binary trees, and then a couple not perfect ones
/////////////////////////////////////////////////////////////////////////////////////

let perfectTree1 = {
	"value": "01",
	"leftChild": {
		"value": 10,    
	},
	"rightChild": {
		"value": 11
  }		
};

let perfectTree2 = {
	"value": "01",
	"leftChild": {
		"value": 10,
    "leftChild": {
				"value": 100
			},
			"rightChild": {
				"value": 101
			}
	},
	"rightChild": {
		"value": 11,
    "leftChild": {
				"value": 110
			},
			"rightChild": {
				"value": 111
			}
	}
};

let bst1 = {
	"value": "f",
	"leftChild": {
		"value": "d",
		"leftChild": {
			"value": "b"
		},
		"rightChild": {
			"value": "e"
		}
	},
	"rightChild": {
		"value": "h"
	}
};

let bst2 = {
	"value": "f",
	"leftChild": {
		"value": "d",
		"leftChild": {
			"value": "b",
			"leftChild": {
				"value": "a"
			},
			"rightChild": {
				"value": "c"
			}
		},
		"rightChild": {
			"value": "e"
		}
	},
	"rightChild": {
		"value": "h",
		"leftChild": {
			"value": "g"
		}
	}
};


/////////////////////////////////////////////////////////////////////////////////////
// 1. STEP-THROUGH VERSION (for easier visualization / drawing)
/////////////////////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////
// SETUP / INITIALIZING VARIBALES
///////////////////////////////////////////////

// Iterate step by step with mouse clicks
document.addEventListener("click", getTreeHeight);

// PARAMS:
let tree = bst2; // test with an example tree from above

// Local vars (if this wasn't a step-through type of loop)
let level = 0; // root will be level 0
let currentSet = {level: 0, currentNode: tree};
let unvisitedNodes = [];

function getTreeHeight () {
  
  // This feels wrong, but my first idea: track the level of each node when added to the stack
  // so : [  {1, node}, {1, node}, {2, node} ... ]
  
  if (currentSet) {
    
    console.log("Current Node " + currentSet.currentNode.value + " -- current level: " + currentSet.level);
    console.log("--------------------------");
    console.log(unvisitedNodes);

    if (currentSet.currentNode.rightChild) {
      unvisitedNodes.push( {level: currentSet.level + 1, currentNode: currentSet.currentNode.rightChild} );
      console.log("Pushed right node " + currentSet.currentNode.rightChild.value);
      
      if (currentSet.level + 1 > level) {
        level = currentSet.level + 1;
        console.log("Increased num of levels to: " + level);
      }
    }
    
    if (currentSet.currentNode.leftChild) {
      unvisitedNodes.push( {level: currentSet.level + 1, currentNode: currentSet.currentNode.leftChild} );
      console.log("Pushed left node " + currentSet.currentNode.leftChild.value);
      
      if (currentSet.level + 1 > level) {
        level = currentSet.level + 1;
        console.log("Increased num of levels to: " + level);
      }
    }
    
//     console.log("Node " + currentSet.currentNode.value);
//     console.log(unvisitedNodes);

    // On next iteration, current node will be the unvisited node
    // that was most recently added to the stack -- last in, first out!
    currentSet = unvisitedNodes.pop();  
    
  } else {
    console.log("HEIGHT OF TREE IS: " + level);
    
    console.log("Sorry, we're fresh outta nodes!"); 
  }
  
  // return level;
  
} // end function


/////////////////////////////////////////////////////////////////////////////////////
// 2. REGULAR VERSION (loop with output, local vars)
/////////////////////////////////////////////////////////////////////////////////////

// Test it with an example tree from above -- output will be height of the tree
console.log ( getTreeHeight(bst2) );

function getTreeHeight (tree) {
  
  // This feels wrong, but my first idea: track the level of each node when added to the stack
  // so : [  {1, node}, {1, node}, {2, node} ... ]
  
  let level = 0; // root will be level 0
  let currentSet = {level: 0, currentNode: tree};
  let unvisitedNodes = [];
  
  while (currentSet) {
    
    console.log("Current Node " + currentSet.currentNode.value + " -- current level: " + currentSet.level);
    console.log("--------------------------");
    console.log(unvisitedNodes);

    if (currentSet.currentNode.rightChild) {
      unvisitedNodes.push( {level: currentSet.level + 1, currentNode: currentSet.currentNode.rightChild} );
      console.log("Pushed right node " + currentSet.currentNode.rightChild.value);
      
      if (currentSet.level + 1 > level) {
        level = currentSet.level + 1;
        console.log("Increased num of levels to: " + level);
      }
    }
    
    if (currentSet.currentNode.leftChild) {
      unvisitedNodes.push( {level: currentSet.level + 1, currentNode: currentSet.currentNode.leftChild} );
      console.log("Pushed left node " + currentSet.currentNode.leftChild.value);
      
      if (currentSet.level + 1 > level) {
        level = currentSet.level + 1;
        console.log("Increased num of levels to: " + level);
      }
    }
    
    // On next iteration, current node will be the unvisited node
    // that was most recently added to the stack -- last in, first out!
    currentSet = unvisitedNodes.pop();  
    
  } // end loop
  
  console.log("HEIGHT OF TREE IS: " + level);
  return level;
  
} // end function

traverseAndDrawDepthFirst.js

// NOTES: - This depends on calling generateTreeWithCoords(), see other file
//        - Using p5js for drawing here
//        - The function traverseAndDrawDepthFirst() runs step-by-step based on mouse clicks

function traverseAndDrawTreeDepthFirst () {
  
  if (currentNode) {
    
    console.log("======================");
    
    console.log("Node " + currentNode.value + " -- level: " + currentNode.level + " -- coords:: " + currentNode.x + ", " + currentNode.y);
    console.log(unvisitedNodes);
    
    // DRAW IT:
    fill( colorArray[currentNode.level - 1] );
    ellipse(currentNode.x, currentNode.y, nodeSize, nodeSize);
    
    if (currentNode.rightChild) {
      unvisitedNodes.push(currentNode.rightChild);
      console.log("Pushed right node " + currentNode.rightChild.value);
    }
    
    if (currentNode.leftChild) {
      unvisitedNodes.push(currentNode.leftChild);
      console.log("Pushed left node " + currentNode.leftChild.value);
    }


    // On next iteration, current node will be the unvisited node
    // that was most recently added to the stack -- last in, first out!
    currentNode = unvisitedNodes.pop();  
    
  } else {
   console.log("Fresh outa nodes!"); 
  }
}


//////////////////////////////////////////
// SETUP / INITIALIZING VARIBALES
/////////////////////////////////////////////////////////////////////////////////////

// Setup for drawing:
const canvasWidth = 1000, canvasHeight = 900;   // canvas dimensions
const nodeSize = 50;                            // given as diameter of each node
const xMargin = 10;                         // minimum X margin between nodes
const yMargin = 25;                         // Y margin between each level

// For drawing each level a different color:
let colorArray = ["#ffffba", "#baffc9", "#bae1ff", "#ffdfba", "#ffb3ba"] ;

////////////////////////////////////////////////////////////////////////////

let myTree = generateTreeWithCoords(5);
console.log("==============================");
console.log( JSON.stringify(myTree, null, '\t' ) ); 

let currentNode = myTree;
let unvisitedNodes = [];
document.addEventListener("click", traverseAndDrawTreeDepthFirst);
console.log("==============================");

////////////////////////////////////////////////////////////////////////////////////////////////////////////
// SETUP FOR P5JS DRAWING:
////////////////////////////////////////////////////////////////////////////////////////////////////////////

function setup() {
  createCanvas(canvasWidth, canvasHeight);

  // 2 frames per second, easier for testing :)
  frameRate(2);

  // DRAW THE TREEEEEE
  //generateAndDrawTree();
  
  noLoop();
}

// Loops for animation, FPS controlled by frameRate() -- default is up to 60 if supported?
function draw() {

} // end draw()

// function mousePressed() {
//   //iterations++;
//   redraw();
// }

tree-traversal.js

/////////////////////////////////////////////////////////////////////////////////////
// EXAMPLE TREES -- for now, easy mode -- perfect binary trees
/////////////////////////////////////////////////////////////////////////////////////

let tree1 = {
	"value": "01",
	"leftChild": {
		"value": 10,    
	},
	"rightChild": {
		"value": 11
  }		
};

let tree2 = {
	"value": "01",
	"leftChild": {
		"value": 10,
    "leftChild": {
				"value": 100
			},
			"rightChild": {
				"value": 101
			}
	},
	"rightChild": {
		"value": 11,
    "leftChild": {
				"value": 110
			},
			"rightChild": {
				"value": 111
			}
	}
};

/////////////////////////////////////////////////////////////////////////////////////
// SETUP / INITIALIZING VARIBALES
/////////////////////////////////////////////////////////////////////////////////////

// Iterate step by step with mouse clicks
// ***** Use any of the functions defined below as the callback here: ********
document.addEventListener("click", traverseTreeBreadthFirst);

// NOTE: To run the code below as a normal loop (instead of stepping through),
// just replace the word "if" with "while" in the code, and you're good to go!

// Initialize with the tree itself as the first node
// ******* Use any of the example trees at the top of this file ********
let currentNode = tree2;
let unvisitedNodes = [];

////////////////////////////////////////////////////////////////////////////

function traverseTreeBreadthFirst () {
  console.log("--------------------------");
  
  if (currentNode) {
    
    console.log("Node " + currentNode.value);
    console.log(unvisitedNodes);
        
    if (currentNode.leftChild) {
      unvisitedNodes.push(currentNode.leftChild);
      console.log("Pushed left node " + currentNode.leftChild.value);
    }
    
    if (currentNode.rightChild) {
      unvisitedNodes.push(currentNode.rightChild);
      console.log("Pushed right node " + currentNode.rightChild.value);
    }
    
    console.log("Node " + currentNode.value);
    console.log(unvisitedNodes);

    // On next iteration, current node will be the unvisited node
    // that was FIRST added to the stack -- first in, first out!
    currentNode = unvisitedNodes.shift();  
    
  } else {
    console.log("Sorry, we're fresh outta nodes!"); 
  }
}



function traverseTreeDepthFirst () {
  console.log("--------------------------");
  
  if (currentNode) {
    
    console.log("Node " + currentNode.value);
    console.log(unvisitedNodes);
    
    if (currentNode.rightChild) {
      unvisitedNodes.push(currentNode.rightChild);
      console.log("Pushed right node " + currentNode.rightChild.value);
    }
    
    if (currentNode.leftChild) {
      unvisitedNodes.push(currentNode.leftChild);
      console.log("Pushed left node " + currentNode.leftChild.value);
    }
    
    console.log("Node " + currentNode.value);
    console.log(unvisitedNodes);

    // On next iteration, current node will be the unvisited node
    // that was most recently added to the stack -- last in, first out!
    currentNode = unvisitedNodes.pop();  
    
  } else {
    console.log("Sorry, we're fresh outta nodes!"); 
  }
}