I have a list of files, an I use the KNN algorithm to classify these files.

dataset = pd.read_csv(file)
training_samples = get_sample_number(dataset)
X_train = dataset.iloc[:training_samples, 5:9]
y_train = dataset.iloc[:training_samples, 9]
X_test = dataset.iloc[training_samples:, 5:9]

# Feature Scaling
sc = StandardScaler()
X_train = sc.fit_transform(X_train)
X_test = sc.fit_transform(X_test)

# Fitting classifier to the training set
classifier = KNeighborsClassifier(n_neighbors=5, metric='minkowski', p=2)
classifier.fit(X_train, y_train)

y_pred = classifier.predict(X_test)

Now I have my categories in my y_pred array. But I want to save the result in the file where I read the dataset. How can I link a prediction to the right row in the file (or dataset)?

  • 1
    $\begingroup$ Nitpicking comment: when scaling your test data, you should use sc.transform not sc.fit_transform. In other words, you should standardize the test data with the mean/std from the training data. $\endgroup$ Commented Jan 18, 2018 at 20:48

2 Answers 2


First as "timleathart" mentioned, you need to fix your code by changing this line :

X_test = sc.fit_transform(X_test)


X_test = sc.transform(X_test)

For your question :

  • you have already the number of samples (training_samples) used for training. so all you need is to iterate over the y_pred and save the values in new column in the dataset starting from "training_samples" as row index.
  • 2
    $\begingroup$ Iterating on pandas objects is generally a bad idea, especially when a simple vectorized alternative exists (as is the case here). In addition to being unidiomatic, it's often way slower than the alternative. $\endgroup$
    – David Marx
    Commented Jan 18, 2018 at 21:42

First and foremost: I would strongly advise against modifying your original data file. It introduces risk that your workflow will become unrepeatable.

To persist your results, I recommend doing something like this:

in_fname = '{}.csv'.format(filename)
out_fname = '{}_SCORED.csv'.format(filename)
dataset = pd.read_csv(in_fname)

... do stuff

dataset.loc[training_samples:, 'scores'] = y_pred
dataset.to_csv(out_fname, header=True, index=False)    

If you really want to overwrite your original data, just set in_fname=out_fname. But I'd advise against it.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.