Added optimized code
Used correct sort_by