from sklearn.model_selection import StratifiedShuffleSplit
sss=StratifiedShuffleSplit(n_splits=1,test_size=0.25,random_state=0)

for train_index, test_index in sss.split(df, df["km_driven_cat"]):
    train_set = df.iloc[train_index]
    test_set = df.iloc[test_index]