Get a "train or test" slicing elegantly after splitting I am learning numpy, pandas, and sklearn for machine learning now. My teacher gives me code like below, but I guess it is ugly. This code is for splitting the data into training set and testing set, and later we could manipulate the entire data or just training/testing set by the [True, False, True, False ... ] idx_train/idx_test list later.

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split

m = 10
x1 = np.arange(m)
x2 = x1 * 2
x3 = x1 * 3
x4 = x1 * 4
x = np.c_[x1, x2, x3, x4]
x = pd.DataFrame(x)

x_train, x_test = train_test_split(x, train_size=0.7, random_state=666)
TRAIN, TEST = 'train', 'test'
TRAIN_TEST = TRAIN + TEST
pd.options.mode.chained_assignment = None
x_train[TRAIN_TEST] = TRAIN
x_test[TRAIN_TEST] = TEST
pd.options.mode.chained_assignment = 'warn'
x = pd.concat([x_train, x_test])
idx_train = x[TRAIN_TEST] == TRAIN
idx_test = np.invert(idx_train)
x_train = x[idx_train]  # later
x_test = x[idx_test]


And I get the below code by myself. I feel it is only a little better.
offset_orig = np.arange(m)
rate = 0.7
offset_train, offset_test = train_test_split(offset_orig, train_size=rate, random_state=666)
if rate &lt; 0.5:
    offset_set_train = set(offset_train)
    idx_train = [i in offset_set_train for i in offset_orig]  # This may be slow. Is numpy broadcast or implicit cycling available here?
    idx_test = np.invert(idx_train)
else:
    offset_set_test = set(offset_test)
    idx_test = [i in offset_set_test for i in offset_orig]  # This may be slow. Is numpy broadcast or implicit cycling available here?
    idx_train = np.invert(idx_test)

x_train = x[idx_train]  # later
x_test = x[idx_test]

I want to get the [True, False, True, False ... ] list idx_train and idx_test for slicing usage later.
I believe that in the world of Python, there should be and will be only one appropriate way to achieve an objective. Could someone give me that very approach to get a &quot;train or test&quot; slicing? (Usage of the lib function train_test_split is required.)
Inspired by Gulzar at Stackoverflow.
I do not realize it works until after a night's sleep.
So, below is enough:
offset_orig = range(m)
offset_train, offset_test = train_test_split(offset_orig, train_size=0.7, random_state=666)

x_train = x.iloc[offset_train]  # later
x_test = x.iloc[offset_test]

The offset_train list ([0, 2, ...]) has the same feature in later programming against the [True, False, True, False ...] list for slicing I pursued at the beginning.
For example, I can use offset_train later as below:

import matplotlib.pyplot as plt
import seaborn as sns
fig, axes = plt.subplots(nrows=2, ncols=2)
fig.set_size_inches(12, 10)
sns.boxplot(data=df.iloc[offset_train], y='count', ax=axes[0][0])
sns.boxplot(data=df.iloc[offset_train], y='count', x='season', ax=axes[0][1])
sns.boxplot(data=df.iloc[offset_train], y='count', x='hour', ax=axes[1][0])
sns.boxplot(data=df.iloc[offset_train], y='count', x='workingday', ax=axes[1][1])


And even as below:

# get noise index
mu_count = df.loc[offset_train, 'count'].mean(axis=0)
sigma_count = df.loc[offset_train, 'count'].std(axis=0)
idx_bad_np = abs(df.loc[offset_train, 'count'] - mu_count) &gt; 3 * sigma_count
idx_good_np = np.invert(idx_bad_np)
# remove noise index
print(f'Before removal: sum(idx good) = {sum(idx_good_np)}, sum(idx bad) = {sum(idx_bad_np)}, len(df[train]) = {len(df.iloc[offset_train])}')
offset_train = list(set(offset_train) - set(np.array(offset_train)[idx_bad_np]))
print(f'After removal: sum(idx good) = {sum(idx_good_np)}, sum(idx bad) = {sum(idx_bad_np)}, len(df[train]) = {len(df.iloc[offset_train])}')
# check again
idx_bad_np = abs(df.loc[offset_train, 'count'] - mu_count) &gt; 3 * sigma_count
print(f'Check sum(idx bad) after removal: {sum(idx_bad_np)}')


Thanks to the community. I feel I have so far to go on the path of Python.