如何返回所有具有相同数量的值0和1的行？

def find_rows_with_equal_zeros_ones(matrix):
    count_dict = {}
    for row in matrix:
        zeros = row.count(0)
        ones = row.count(1)
        diff = abs(zeros - ones)
        if diff in count_dict:
            count_dict[diff].append(row)
        else:
            count_dict[diff] = [row]
    
    result = []
    for rows in count_dict.values():
        result.extend(rows)
    
    return result