def test_balance_tree(self): node_1, node_3 = create_unbalanced_tree_with_zig_zag() tree = AVL(node_1) tree._balance(node_3) node_1 = tree.get_root() node_0 = node_1.get_left_child() node_3 = node_1.get_right_child() node_2 = node_3.get_left_child() node_4 = node_3.get_right_child() expect(node_0.get_key()).to.be.equal(0) expect(node_1.get_key()).to.be.equal(1) expect(node_2.get_key()).to.be.equal(2) expect(node_3.get_key()).to.be.equal(3) expect(node_4.get_key()).to.be.equal(4) expect(node_0.get_height()).to.be.equal(0) expect(node_1.get_height()).to.be.equal(2) expect(node_2.get_height()).to.be.equal(0) expect(node_3.get_height()).to.be.equal(1) expect(node_4.get_height()).to.be.equal(0) expect(node_0.get_node_count()).to.be.equal(1) expect(node_1.get_node_count()).to.be.equal(5) expect(node_2.get_node_count()).to.be.equal(1) expect(node_3.get_node_count()).to.be.equal(3) expect(node_4.get_node_count()).to.be.equal(1)
def test_balance_tree_with_same_weight_on_second_level(self): node_5 = create_unbalanced_tree_with_equal_weight_on_second_level() tree = AVL(node_5) tree._balance(node_5) node_4 = tree.get_root() node_2 = node_4.get_left_child() node_5 = node_4.get_right_child() node_3 = node_5.get_left_child() expect(node_4.get_key()).to.equal(4) expect(node_2.get_key()).to.equal(2) expect(node_5.get_key()).to.equal(5) expect(node_3.get_key()).to.equal(3)