Moved diagonal back outside of the conditional while. The sum of the rows in the affinity matrix never change, so this only needs to be calculated once.