41 lines
1.0 KiB
Plaintext
41 lines
1.0 KiB
Plaintext
// Define ColoredTrees with red and black nodes and integer leaves
|
|
data ColoredTree = leaf(int N)
|
|
| red(ColoredTree left, ColoredTree right)
|
|
| black(ColoredTree left, ColoredTree right);
|
|
|
|
// Count the number of black nodes
|
|
public int cntBlack(ColoredTree t){
|
|
int c = 0;
|
|
visit(t) {
|
|
case black(_,_): c += 1;
|
|
};
|
|
return c;
|
|
}
|
|
|
|
// Returns if a tree is balanced
|
|
public bool balance(ColoredTree t){
|
|
visit(t){
|
|
case black(a,b): if (cntBlack(a) == cntBlack(b)) true; else return false;
|
|
case red(a,b): if (cntBlack(a) == cntBlack(b)) true; else return false;
|
|
}
|
|
return true;
|
|
}
|
|
// Compute the sum of all integer leaves
|
|
public int addLeaves(ColoredTree t){
|
|
int c = 0;
|
|
visit(t) {
|
|
case leaf(int N): c += N;
|
|
};
|
|
return c;
|
|
}
|
|
|
|
// Add green nodes to ColoredTree
|
|
data ColoredTree = green(ColoredTree left, ColoredTree right);
|
|
|
|
// Transform red nodes into green nodes
|
|
public ColoredTree makeGreen(ColoredTree t){
|
|
return visit(t) {
|
|
case red(l, r) => green(l, r)
|
|
};
|
|
}
|