Added stratify option to train_test_split function.
This PR adds `stratify` option to `train_test_split` method. I took reference from scikit-learn's `StratifiedShuffleSplit` class for implementing stratified split and integrated the changes as were suggested by @lhoestq.

It fixes #3452.

@lhoestq Please review and let me know, if any changes are required.
