def solution(data, n): memo = {} for i in data: if i not in memo: memo[i] = 1 else: memo[i] += 1 return [d for d in data if memo[d] <= n] tests = [ ([[1, 2, 3], 0], []), ([[1, 2, 2, 3, 3, 3, 4, 5, 5], 1], [1, 4]), ] for i, o in tests: result = solution(*i) print (i, result == o, result, o)