Cleanup the code and O(|V| + |E|)