Google Foobar Level 1 - Minion Task Schedule

Question

I recently just stumbled upon a post regarding Google Foobar Level 2 and decided to try it out myself but starting from Level 1. The description and code are shown below.

Write a function called answer(data, n) that takes in a list of less than 100 integers and a number n, and returns that same list but with all of the numbers that occur more than n times removed entirely. The returned list should retain the same ordering as the original list - you don't want to mix up those carefully planned shift rotations! For instance, if data was [5, 10, 15, 10, 7] and n was 1, answer(data, n) would return the list [5, 15, 7] because 10 occurs twice, and was thus removed from the list entirely.

Constrains:

Standard libraries are supported except for bz2, crypt, fcntl, mmap, pwd, pyexpat, select, signal, termios, thread, time, unicodedata, zipimport, zlib.

Test Cases:

Inputs:
(int list) data = [1, 2, 3]
(int) n = 0
Output:
(int list) []

Inputs:
(int list) data = [1, 2, 2, 3, 3, 3, 4, 5, 5]
(int) n = 1
Output:
(int list) [1, 4]

Inputs:
(int list) data = [1, 2, 3]
(int) n = 6
Output:
(int list) [1, 2, 3]

Code:

def answer(data, n):
    if len(data) > 99:
        exit('List contains more than 100 elements')
    print('data:\n', data) #for error-checking

    count = dict()
    for i in data:
        count.setdefault(i, 0)
        count[i] += 1
    print('count:\n', count) #for error-checking

    for k, v in count.items():
        if v > n:
            for i in range(v):
                data.remove(k)

    return data

I decided to also test the code myself by using test cases and a new list containing random numbers

data1 = [1, 2, 3]
n1 = 0

data2 = [1, 2, 2, 3, 3, 3, 4, 5, 5]
n2 = 1

data3 = [1, 2, 3]
n3 = 6

data4 = [random.randint(1, 15) for i in range(0, 99)]
n4 = 6

ans1 = answer(data1, n1)
print('ans1:\n', ans1, '\n')

ans2 = answer(data2, n2)
print('ans2:\n', ans2, '\n')

ans3 = answer(data3, n3)
print('ans3:\n', ans3, '\n')

ans4 = answer(data4, n4)
print('ans4:\n', ans4, '\n')

Output:

data:
 [1, 2, 3]
count:
 {1: 1, 2: 1, 3: 1}
ans1:
 []

data:
 [1, 2, 2, 3, 3, 3, 4, 5, 5]
count:
 {1: 1, 2: 2, 3: 3, 4: 1, 5: 2}
ans2:
 [1, 4]

data:
 [1, 2, 3]
count:
 {1: 1, 2: 1, 3: 1}
ans3:
 [1, 2, 3]

data:
 [15, 9, 14, 10, 13, 8, 7, 2, 10, 3, 10, 11, 9, 9, 10, 7, 5, 11, 12, 14, 1, 5, 14, 15, 11, 2, 13, 2, 6, 1, 6, 8, 3, 15, 8, 6, 6, 13, 14, 14, 2, 14, 12, 6, 15, 4, 7, 5, 13, 8, 8, 4, 7, 15, 6, 1, 2, 11, 14, 6, 8, 4, 10, 4, 2, 9, 6, 14, 11, 12, 8, 15, 7, 15, 6, 1, 14, 8, 5, 14, 6, 10, 8, 1, 4, 15, 11, 9, 11, 10, 13, 13, 9, 7, 7, 4, 4, 12, 12]
count:
 {15: 8, 9: 6, 14: 10, 10: 7, 13: 6, 8: 9, 7: 7, 2: 6, 3: 2, 11: 7, 5: 4, 12: 5, 1: 5, 6: 10, 4: 7}
ans4:
 [9, 13, 2, 3, 9, 9, 5, 12, 1, 5, 2, 13, 2, 1, 3, 13, 2, 12, 5, 13, 1, 2, 2, 9, 12, 1, 5, 1, 9, 13, 13, 9, 12, 12]

Any input is appreciated. Thank you :-)

Answer 1

if len(data) > 99:
    exit('List contains more than 100 elements')

I'd change that error message to something like 'List contains a hundred or more elements'. I might also raise a ValueError instead of calling exit. I think that's more appropriate for an error that the user may want to catch.

---

You can get rid of your call to setdefault by using a defaultdict instead:

from collections import defaultdict

. . .

count = defaultdict(int)  # int is a function that returns 0 when no args are given to it
for i in data:
    count[i] += 1

The gain here isn't huge, but they're handy in many cases.

collections.Counter would probably be appropriate here as well, but I've only used it once, so I can't show a good example of it.

---