Slight update to count_matrix.dia sketch