Instead of `tri_idx`, I think you can use `np.unravel_index`. May be faster and uses less memory.