Merge pull request #66 from clockor2/dev

feat: add `Node.applyPostOrder` method and tests
clockor2 · May 22, 2024 · 09dd8b8 · 09dd8b8
2 parents 5a80f9e + b481c0d
commit 09dd8b8
Show file tree

Hide file tree

Showing 2 changed files with 181 additions and 97 deletions.
diff --git a/src/tree/node.ts b/src/tree/node.ts
@@ -146,8 +146,8 @@ export class Node {
   }
 
   /**
-   * Apply a function `f()` to each node in a subtree descending from `node` in post-order
-   * @param {Node} node Node from which to apply `f()` post-order
+   * Apply a function `f()` to each node in a subtree descending from `node` in pre-order
+   * @param {Node} node Node from which to apply `f()` pre-order
    */
   applyPreOrder<T>(f: (node: Node) => T): T[] {
     const res: T[] = [];
@@ -168,4 +168,41 @@ export class Node {
     }
     return res;
   }
+
+  /**
+   * Apply a function `f()` to each node in a subtree descending from `node` in post-order
+   * @param {Node} node Node from which to apply `f()` post-order
+   */
+  applyPostOrder<T>(f: (node: Node) => T): T[] {
+    const res: T[] = [];
+    const stack: Node[] = [this]; // Initialize the stack with the root node.
+    const visited: Set<Node> = new Set(); // Track visited nodes.
+
+    while (stack.length > 0) {
+      const currentNode = stack[stack.length - 1]; // Peek the last node from the stack.
+
+      if (!currentNode) {
+        stack.pop(); // If the node is null or undefined, remove it from the stack.
+        continue;
+      }
+
+      let allChildrenVisited = true;
+      for (let i = currentNode.children.length - 1; i >= 0; i--) {
+        if (!visited.has(currentNode.children[i])) {
+          stack.push(currentNode.children[i]); // Push unvisited children to the stack.
+          allChildrenVisited = false;
+        }
+      }
+
+      if (allChildrenVisited) {
+        const lastNode = stack.pop(); // Pop the last node from the stack.
+        if (lastNode) {
+          visited.add(lastNode); // Mark the node as visited.
+          const thisRes = f(lastNode);
+          res.push(thisRes); // Apply f to the current node and collect the result.
+        }
+      }
+    }
+    return res;
+  }
 }
diff --git a/test/tree/node.spec.ts b/test/tree/node.spec.ts
@@ -1,131 +1,178 @@
 import { Node } from '../../src';
 
 describe('Node', () => {
-  test('create node and check basic properties', () => {
-    const node = new Node(1);
-    expect(node.id).toBe(1);
-    expect(node.parent).toBeUndefined();
-    expect(node.children.length).toBe(0);
-  });
+    test('create node and check basic properties', () => {
+        const node = new Node(1);
+        expect(node.id).toBe(1);
+        expect(node.parent).toBeUndefined();
+        expect(node.children.length).toBe(0);
+    });
 
-  test('addChild and isLeaf', () => {
-    const parentNode = new Node(1);
-    const childNode = new Node(2);
+    test('addChild and isLeaf', () => {
+        const parentNode = new Node(1);
+        const childNode = new Node(2);
 
-    parentNode.addChild(childNode);
+        parentNode.addChild(childNode);
 
-    expect(parentNode.children.length).toBe(1);
-    expect(childNode.parent).toBe(parentNode);
-    expect(parentNode.isLeaf()).toBe(false);
-    expect(childNode.isLeaf()).toBe(true);
-  });
+        expect(parentNode.children.length).toBe(1);
+        expect(childNode.parent).toBe(parentNode);
+        expect(parentNode.isLeaf()).toBe(false);
+        expect(childNode.isLeaf()).toBe(true);
+    });
 
-  test('removeChild', () => {
-    const parentNode = new Node(1);
-    const childNode = new Node(2);
+    test('removeChild', () => {
+        const parentNode = new Node(1);
+        const childNode = new Node(2);
 
-    parentNode.addChild(childNode);
-    parentNode.removeChild(childNode);
+        parentNode.addChild(childNode);
+        parentNode.removeChild(childNode);
 
-    expect(parentNode.children.length).toBe(0);
-    expect(childNode.parent).toBeUndefined();
-  });
+        expect(parentNode.children.length).toBe(0);
+        expect(childNode.parent).toBeUndefined();
+    });
 
-  test('isRoot', () => {
-    const parentNode = new Node(1);
-    const childNode = new Node(2);
+    test('isRoot', () => {
+        const parentNode = new Node(1);
+        const childNode = new Node(2);
 
-    parentNode.addChild(childNode);
+        parentNode.addChild(childNode);
 
-    expect(parentNode.isRoot()).toBe(true);
-    expect(childNode.isRoot()).toBe(false);
-  });
+        expect(parentNode.isRoot()).toBe(true);
+        expect(childNode.isRoot()).toBe(false);
+    });
 
-  test('isSingleton', () => {
-    const parentNode = new Node(1);
-    const childNode1 = new Node(2);
-    const childNode2 = new Node(3);
+    test('isSingleton', () => {
+        const parentNode = new Node(1);
+        const childNode1 = new Node(2);
+        const childNode2 = new Node(3);
 
-    parentNode.addChild(childNode1);
+        parentNode.addChild(childNode1);
 
-    expect(parentNode.isSingleton()).toBe(true);
-    expect(childNode1.isSingleton()).toBe(false); // childNode1 is a leaf, not a singleton
+        expect(parentNode.isSingleton()).toBe(true);
+        expect(childNode1.isSingleton()).toBe(false); // childNode1 is a leaf, not a singleton
 
-    parentNode.addChild(childNode2);
+        parentNode.addChild(childNode2);
 
-    expect(parentNode.isSingleton()).toBe(false);
-    expect(childNode1.isSingleton()).toBe(false);
-  });
+        expect(parentNode.isSingleton()).toBe(false);
+        expect(childNode1.isSingleton()).toBe(false);
+    });
 
-  test('isHybrid', () => {
-    const node = new Node(1);
+    test('isHybrid', () => {
+        const node = new Node(1);
 
-    expect(node.isHybrid()).toBe(false);
+        expect(node.isHybrid()).toBe(false);
+
+        node.hybridID = 42;
 
-    node.hybridID = 42;
+        expect(node.isHybrid()).toBe(true);
+    });
 
-    expect(node.isHybrid()).toBe(true);
-  });
+    test('getAncestors', () => {
+        const node1 = new Node(1);
+        const node2 = new Node(2);
+        const node3 = new Node(3);
 
-  test('getAncestors', () => {
-    const node1 = new Node(1);
-    const node2 = new Node(2);
-    const node3 = new Node(3);
+        node1.addChild(node2);
+        node2.addChild(node3);
 
-    node1.addChild(node2);
-    node2.addChild(node3);
+        expect(node1.getAncestors()).toEqual([node1]);
+        expect(node2.getAncestors()).toEqual([node2, node1]);
+        expect(node3.getAncestors()).toEqual([node3, node2, node1]);
+    });
 
-    expect(node1.getAncestors()).toEqual([node1]);
-    expect(node2.getAncestors()).toEqual([node2, node1]);
-    expect(node3.getAncestors()).toEqual([node3, node2, node1]);
-  });
+    test('copy', () => {
+        const node1 = new Node(1);
+        const node2 = new Node(2);
+        const node3 = new Node(3);
 
-  test('copy', () => {
-    const node1 = new Node(1);
-    const node2 = new Node(2);
-    const node3 = new Node(3);
+        node1.addChild(node2);
+        node2.addChild(node3);
 
-    node1.addChild(node2);
-    node2.addChild(node3);
+        const copyNode1 = node1.copy();
 
-    const copyNode1 = node1.copy();
+        expect(copyNode1).not.toBe(node1);
+        expect(copyNode1).toEqual(node1);
+    });
 
-    expect(copyNode1).not.toBe(node1);
-    expect(copyNode1).toEqual(node1);
-  });
+    test('toString', () => {
+        const node = new Node(1);
+        expect(node.toString()).toBe('node#1');
+    });
 
-  test('toString', () => {
-    const node = new Node(1);
-    expect(node.toString()).toBe('node#1');
-  });
+    test('isLeftOf', () => {
+        const nodeA = new Node(1);
+        const nodeB = new Node(2);
+        const nodeC = new Node(3);
+        const nodeD = new Node(4);
 
-  test('isLeftOf', () => {
-    const nodeA = new Node(1);
-    const nodeB = new Node(2);
-    const nodeC = new Node(3);
-    const nodeD = new Node(4);
+        nodeA.addChild(nodeB);
+        nodeA.addChild(nodeC);
+        nodeC.addChild(nodeD);
 
-    nodeA.addChild(nodeB);
-    nodeA.addChild(nodeC);
-    nodeC.addChild(nodeD);
+        expect(nodeB.isLeftOf(nodeD)).toBe(true);
+        expect(nodeD.isLeftOf(nodeB)).toBe(false);
+        expect(nodeA.isLeftOf(nodeD)).toBeUndefined();
+    });
 
-    expect(nodeB.isLeftOf(nodeD)).toBe(true);
-    expect(nodeD.isLeftOf(nodeB)).toBe(false);
-    expect(nodeA.isLeftOf(nodeD)).toBeUndefined();
-  });
+    test('applyPreOrder', () => {
+        const nodeA = new Node(1);
+        const nodeB = new Node(2);
+        const nodeC = new Node(3);
 
-  test('applyPreOrder', () => {
-    const nodeA = new Node(1);
-    const nodeB = new Node(2);
-    const nodeC = new Node(3);
+        nodeA.addChild(nodeB);
+        nodeA.addChild(nodeC);
 
-    nodeA.addChild(nodeB);
-    nodeA.addChild(nodeC);
+        const nodes = nodeA.applyPreOrder(node => node);
+        expect(nodes.length).toBe(3);
+        expect(nodes).toContain(nodeA);
+        expect(nodes).toContain(nodeB);
+        expect(nodes).toContain(nodeC);
+    });
+
+    test('applyPostOrder', () => {
+        const nodeA = new Node(1);
+        const nodeB = new Node(2);
+        const nodeC = new Node(3);
+
+        nodeA.addChild(nodeB);
+        nodeA.addChild(nodeC);
+
+        const nodes = nodeA.applyPostOrder(node => node);
+        expect(nodes.length).toBe(3);
+        expect(nodes).toContain(nodeB);
+        expect(nodes).toContain(nodeC);
+        expect(nodes).toContain(nodeA);
+    });
+
+    test('applyPostOrder to ger root age', () => {
+        const nodeA = new Node(1);
+        const nodeB = new Node(2);
+        const nodeC = new Node(3);
+
+        nodeA.addChild(nodeB);
+        nodeA.addChild(nodeC);
+
+        const ages = nodeA.applyPostOrder(node => {
+            let res;
+            let childBranchLength;
+
+            if (node.isLeaf()) {
+                childBranchLength = 0
+            }
+            else {
+                childBranchLength = 1 // set default branch length to 1
+            }
+
+            node.branchLength !== undefined
+                ? res = childBranchLength + node.branchLength
+                : res = childBranchLength
+
+            return res;
+
+        });
+
+        expect(ages.length).toBe(3);
+        expect(ages[2]).toBe(1);
+    });
 
-    const nodes = nodeA.applyPreOrder(node => node);
-    expect(nodes.length).toBe(3);
-    expect(nodes).toContain(nodeA);
-    expect(nodes).toContain(nodeB);
-    expect(nodes).toContain(nodeC);
-  });
 });