added code that colour the tree branches correctly