Holy hell

x = train_data[train_data[:,1].argsort()][(lambda label_count: list(itertools.chain.from_iterable(range(i, i+50) for i in [0]+label_count[:-1]))+list(itertools.chain.from_iterable(range(i+50, j) for i,j in zip([0]+label_count, label_count))))(list(itertools.accumulate(labels.count(a) for a in range(500))))]