how to impute missing values on numpy array created by train_test_split from pandas.DataFrame?
Can you just cast your np.array from train_test_split back into a pandas.DataFrame so you can carry out your same strategy. This is very common to what I do when dealing with pandas and scikit. For example,
 a = train_test_split
 new_df = pd.DataFrame(a)