Improve performance on the trans_matrix construction

It looks like constructing the trans_matrix is slow.