Updated probability class calculation