import torch


def get_top_n_boxes(box_list, num):
    """
    Args:
        box_list: ([N_i,5],[N_,],[N, C]) x T*1
        num: number of box for each in T*1

    Returns:
        1. [[N_j, 5] x T*1]
        2. [[N_j, C] x T*1]
        3. [[N_j,] x T*1]
         where N_j <= num
    """
    good_box_list = []
    feat_seq = []
    ind_list = []
    for boxes, labels, feat in box_list:
        num_boxes = min(num, len(boxes))
        _, ind = boxes[:, -1].topk(num_boxes)
        good_box_list.append(boxes[ind])
        feat_seq.append(feat[ind])
        ind_list.append(ind)
    feat_seq = torch.cat(feat_seq, 0)

    return good_box_list, feat_seq, ind_list