데이터 세트를 나누다 보면 특정 분류는 학습 세트에는 많고, 시험 세트에는 너무 적어 균형 있게 학습하지 못하는 현상이 발생하기도 한다. 이를 방지하기 위해 train_test_split()으로 데이터 세트를 나눌 때 stratify에도 정답값을 지정해 주었다. stratify에 정답을 지정하면 학습 세트와 시험 세트의 정답 비율을 맞춰서 나눠 준다. 다음 코드를 실행해 결과를 보면 경제 분류의 학습/시험 세트 정답 비율이 비슷한 비율(0.38)로 나뉜 것을 확인할 수 있다.

display(y_train.mean())
display(y_test.mean())

실행 결과

경제   0.384795
복지   0.101754
행정   0.513450
dtype: float64
경제   0.385514
복지   0.100467
행정   0.514019
dtype: float64