Fix the per-fold feature selector for genomics data

Review the AI's codeMedium15 min

Your task

An AI selected the 50 best of 20,000 genomics features with SelectKBest on the full dataset, then cross-validated a classifier on them. The selector had already seen every test fold's labels, so the cross-validated accuracy of 0.97 was a mirage. The fix reruns the selection inside each fold, on that fold's training samples only.

An assistant wrote select_features_per_fold(feature_rows, labels, fold_ids, top_k). It is in the editor. It has a bug. Find it and fix it.

Input and output:

  • feature_rows: one list of feature values per sample.
  • labels: each sample's label, at the same index. 1 is positive, anything else is negative.
  • fold_ids: each sample's fold id, at the same index. Fold ids are whole numbers, not necessarily 0, 1, 2 and so on.
  • top_k: how many features to select in each fold.
  • Return one list of selected feature indexes per fold.

Rules the tests check:

  • Use only the first n samples, where n is the length of the shortest list. If n is 0, return an empty list.
  • The number of features is the length of the first sample's row. A shorter row counts its missing values as 0, and values past that length are ignored.
  • Go through the distinct fold ids in ascending order, one output list each.
  • A fold's training samples are all the samples with a different fold id. A feature's score is the absolute value of its sum over positive training samples minus its sum over negative training samples.
  • Select the top_k features with the highest scores, breaking ties by the lower feature index. Select none if top_k is 0 or less, and all of them if it is more than the number of features.
  • Each fold's list is sorted ascending.

Examples

select_features_per_fold(feature_rows: list[list[int]], labels: list[int], fold_ids: list[int], top_k: int) -> list[list[int]]
Example 1
Inputselect_features_per_fold([[1, 0, 7], [1, 0, 0], [0, 9, 0], [0, 9, 0]], [1, 1, 0, 0], [0, 1, 0, 1], 2)
Output[[0, 1], [1, 2]]
Example 2
Inputselect_features_per_fold([[3, 1], [0, 4], [2, 0]], [1, 0, 1], [0, 1, 1], 1)
Output[[1], [0]]