Bird
Raised Fist0

Implement Train-Test Split with Stratification

CODINGbeginnerInterview freq: high
Choose your preparation mode6 modes available
LearnTraceDebugTrade-offChallenge
⚖️
Train-Test Split with Stratification vs Simple Train-Test Split

When preparing a dataset for a classification task, would you use a stratified train-test split or a simple train-test split? Justify your choice.

📊
Comparison
DimensionTrain-Test Split with StratificationSimple Train-Test Split
Implementation ComplexityMediumLow
Class Distribution MaintenanceHighLow
Suitability for Imbalanced DataHighLow
Computational EfficiencyMediumHigh
Ease of UseMediumHigh
✅
Verdict
Default to stratified train-test split for datasets with significant class imbalance and n > 1000. Use simple train-test split for small, balanced datasets or when computational resources are limited.
💡
Key Insight
Stratified splits ensure class distribution is maintained, which is critical for imbalanced datasets.
⚠️
Interview Trap
❌ Wrong: Always use stratified split because it maintains class distribution.
Ignoring the computational overhead and complexity for small or balanced datasets.
✅ Consider the dataset size and class distribution. Use stratified split for large, imbalanced datasets where class distribution is crucial for evaluation.
Train-Test Split with Stratification
"Maintains class distribution across train and test sets."
Weaknesses
  • Slightly more complex to implement than a simple split.
  • May not work well with very small datasets where class distribution is already skewed.
  • Requires knowledge of the target variable during the split process.
  • Can be computationally more intensive for very large datasets.
Ideal When
  • When n > 1000 and class imbalance is significant.
  • When the dataset has more than 2 classes.
  • When maintaining class distribution is critical for model evaluation.
  • When the dataset is used for a classification task.
Kaggle competitions often use stratified splits to ensure fair evaluation across imbalanced classes.
Simple Train-Test Split
"Quick and easy to implement."
Weaknesses
  • Does not maintain class distribution, leading to potential bias in evaluation.
  • Not suitable for imbalanced datasets.
  • May result in overfitting if the test set is not representative.
  • Random splits can lead to high variance in model performance.
Ideal When
  • When n < 1000 and class distribution is uniform.
  • When the dataset is balanced across classes.
  • When quick prototyping is needed without concern for class distribution.
  • When computational resources are limited.
In early-stage startups, quick prototyping often uses simple splits to iterate rapidly on model ideas.
🐍
Train-Test Split with Stratification - Example
▼ Show
Python
import numpy as np

def stratified_train_test_split(X, y, test_size=0.2):
    unique_classes, y_indices = np.unique(y, return_inverse=True)
    train_indices = []
    test_indices = []
    for cls in unique_classes:
        cls_indices = np.where(y_indices == cls)[0]
        np.random.shuffle(cls_indices)
        split_idx = int(len(cls_indices) * (1 - test_size))
        train_indices.extend(cls_indices[:split_idx])
        test_indices.extend(cls_indices[split_idx:])
    return X[train_indices], X[test_indices], y[train_indices], y[test_indices]
🐍
Simple Train-Test Split - Example
▼ Show
Python
import numpy as np

def simple_train_test_split(X, y, test_size=0.2):
    indices = np.arange(X.shape[0])
    np.random.shuffle(indices)
    split_idx = int(len(indices) * (1 - test_size))
    train_indices = indices[:split_idx]
    test_indices = indices[split_idx:]
    return X[train_indices], X[test_indices], y[train_indices], y[test_indices]
✅
Use Train-Test Split with Stratification When
  • •When n > 1000 and class imbalance is significant.
  • •When the dataset has more than 2 classes.
  • •When maintaining class distribution is critical for model evaluation.
✅
Use Simple Train-Test Split When
  • •When n < 1000 and class distribution is uniform.
  • •When quick prototyping is needed without concern for class distribution.
  • •When computational resources are limited.
↳
Nuances
  • For very small datasets, stratification might not be beneficial as the class distribution is already skewed.
  • In cases where computational resources are extremely limited, a simple split might be the only feasible option.
  • Stratification is crucial when the evaluation metric is sensitive to class imbalance, such as F1-score.

Practice

(1/5)
1. What is the time complexity of performing K-Fold cross-validation on a dataset with n samples and k folds?
easy
A. O(n^2)
B. O(n)
C. O(kn)
D. O(k^2n)

Solution

  1. Step 1: Analyze K-Fold process

    Each fold requires training on (k-1)/k of the data and testing on 1/k.

  2. Step 2: Eliminate wrong options

    O(n) ignores the k factor. O(n^2) and D overestimate complexity.

  3. Final Answer:

    Option C -> Option C

  4. Quick Check:

    K-Fold complexity is O(kn). [OK]
Hint: K-Fold complexity is O(kn). [OK]
Common Mistakes:
2. What does the following Python code output when computing the ROC curve? 
from sklearn.metrics import roc_curve
import numpy as np

y_true = np.array([0, 0, 1, 1])
y_scores = np.array([0.1, 0.4, 0.35, 0.8])
fpr, tpr, thresholds = roc_curve(y_true, y_scores)
print(fpr, tpr, thresholds)
easy
A. array([0. , 0.5, 0.5, 1. ]), array([0. , 0. , 1. , 1. ]), array([1.8, 0.8, 0.4, 0.1])
B. array([0. , 0. , 0.5, 1. ]), array([0. , 0.5, 1. , 1. ]), array([0.8, 0.4, 0.35, 0.1])
C. array([0. , 0.5, 1. ]), array([0. , 0.5, 1. ]), array([0.8, 0.35, 0.1])
D. array([0. , 0.5, 1. ]), array([0. , 1. , 1. ]), array([0.8, 0.35, 0.1])

Solution

  1. Step 1: Understand ROC curve output

    The ROC curve function returns false positive rates (fpr), true positive rates (tpr), and thresholds.

  2. Step 2: Analyze the code

    Given the scores and true labels, the ROC curve is calculated correctly, resulting in fpr, tpr, and thresholds as in array([0. , 0.5, 1. ]), array([0. , 0.5, 1. ]), array([0.8, 0.35, 0.1]).

  3. Final Answer:

    Option C -> Option C

  4. Quick Check:

    ROC curve outputs fpr, tpr, thresholds [OK]
Hint: ROC curve outputs fpr, tpr, thresholds [OK]
Common Mistakes:
3. In a pipeline that includes flipping, cropping, and normalizing images, where should normalization be applied for best results?
medium
A. Before flipping and cropping
B. After flipping but before cropping
C. After cropping but before flipping
D. After both flipping and cropping

Solution

  1. Step 1: Understand the impact of operations

    Normalization should be the last step to ensure consistent scaling.

  2. Step 2: Evaluate each option

    Applying normalization last ensures it accounts for all transformations.

  3. Final Answer:

    Option D -> Option D

  4. Quick Check:

    Normalize last for consistent scaling. [OK]
Hint: Normalize last for consistent scaling. [OK]
Common Mistakes:
4. In a custom dropout layer implementation, how should you modify the forward pass for inference mode to ensure correct scaling?
medium
A. Divide inputs by (1 - dropout_rate).
B. Multiply inputs by (1 - dropout_rate).
C. Apply the same dropout mask as in training.
D. Use a different dropout rate for inference.

Solution

  1. Step 1: Understand Inference Scaling

    During inference, inputs should be scaled by dividing by (1 - dropout_rate) to maintain expected values.

  2. Step 2: Eliminate Wrong Options

    Apply the same dropout mask as in training. incorrectly applies dropout during inference. Multiply inputs by (1 - dropout_rate). scales incorrectly. Use a different dropout rate for inference. suggests changing dropout rate, which is not standard practice.

  3. Final Answer:

    Option A -> Option A

  4. Quick Check:

    Scale by dividing during inference. [OK]
Hint: Scale by dividing during inference. [OK]
Common Mistakes:
5. In your implementation of K-Means, the algorithm is not converging. What is the most likely cause?
medium
A. The dataset is too small.
B. The number of clusters is too large.
C. The initial centroids are poorly chosen.
D. The distance metric is incorrect.

Solution

  1. Step 1: Identify convergence issues

    Poorly chosen initial centroids can lead to non-convergence as they may not represent the data well.

  2. Step 2: Eliminate wrong options

    The number of clusters is too large. is incorrect because a large number of clusters doesn't prevent convergence. The dataset is too small. is incorrect as a small dataset doesn't inherently cause non-convergence. The distance metric is incorrect. is incorrect if the metric is consistent.

  3. Final Answer:

    Option C -> Option C

  4. Quick Check:

    Check initial centroids for convergence issues. [OK]
Hint: Check initial centroids for convergence issues. [OK]
Common Mistakes:
Recommended

Ready to prove it?

Pass the test and this pattern moves unknown strong while earning +1 Grip toward your certificate.

10 questions~5 mineasy hard