Problem
Given an array of size n that has the following specifications:
Each element in the array contains either a policeman or a thief.
Each policeman can catch only one thief.
A policeman cannot catch a thief who is more than K units away from the policeman.
We need to find the maximum number of thieves that can be caught.
Here’s code for policemen catch thieves written in Python3, I would like feedback/suggestions.
import random
def tp_build(n):
thief_police = []
for i in range(n):
thief_police.append(random.choice(['T', 'P']))
return thief_police
def catch(k=2, n=10):
tp_list = tp_build(n)
print('List: %s' % tp_list)
if 'P' not in tp_list:
return 'No Police!'
if 'T' not in tp_list:
return 'No Thieves!'
p_indexes = []
t_indexes = []
for i in range(n):
if tp_list[i] == 'P':
p_indexes.append(i)
elif tp_list[i] == 'T':
t_indexes.append(i)
combinations = []
for limit in range(1, k + 1):
for police in p_indexes:
if police + limit in t_indexes:
combinations.append((police, police + limit))
for police in p_indexes:
if police - limit in t_indexes:
combinations.append((police, police - limit))
p_list = []
t_list = []
for i, j in combinations:
p_list.append(i)
t_list.append(j)
new_p = []
new_t = []
for i in p_list:
if i not in new_p:
new_p.append(i)
for j in t_list:
if j not in new_t:
new_t.append(j)
final_combinations = list(zip(new_p, new_t))
print('Number of thieves caught: %s'%(len(final_combinations)))
return len(final_combinations)
if __name__ == '__main__':
for i in range(100):
catch()
Solution
Bugs
You return 'No Police!'
instead of returning 0, but you don’t check the return value. You might consider printing the string and returning 0 instead. Likewise for the no-thieves case.
Style
You are nominally compliant with pep-8, but your style is very much a beginner style. Why do you have variables named t_list
and p_list
? Do you really think that the _list
part of the name is the most important, and likely to be overlooked by readers?
When choosing a name, you should almost always avoid using the type of the data in the name (1). After all, the type might change! (For example, you might switch from a list to a dict for storage.)
(1): Except when the type is the important part, like if you’re writing a type conversion function (str2int or wav2mp3 or something).
Iteration
You are missing some tricks of Python iteration. First, let’s look at your tp_build
function:
def tp_build(n):
thief_police = []
for i in range(n):
thief_police.append(random.choice(['T', 'P']))
return thief_police
This code is a one-off “helper” that does a job that is described in the problem statement: generate a random array of police & thieves. Because it’s simple and the requirements are clearly explained by the problem statement, you don’t need to worry about documenting this, or leaving the code spelled out for comprehensibility. So just use a list comprehension and shorten that code a lot (also, strings are iterable):
def tp_build(n):
return [random.choice("TP") for _ in range(n)]
In your “catch` function, you have a lot of iteration. Again, you’re writing too-simple loops when you could get clarity and performance from Python features.
p_indexes = []
t_indexes = []
for i in range(n):
if tp_list[i] == 'P':
p_indexes.append(i)
elif tp_list[i] == 'T':
t_indexes.append(i)
In this paragraph, you construct lists of the indexes of the different characters. If you’re going to look up the value of tp_list[i]
, you should just use enumerate
:
p_indexes = []
t_indexes = []
for i, ch in enumerate(tp_list):
if ch == 'P':
p_indexes.append(i)
else:
t_indexes.append(i)
Of course, this is another case where you could use comprehensions. In theory, this is twice as slow. In practice, the comprehension might be faster because Python knows how to speed it up. You’ll have to check it:
p_indexes = [i for i in range(n) if tp_list[i] == 'P']
t_indexes = [i for i in range(n) if tp_list[i] == 'T']
Clarity
You don’t provide a function level comment explaining your algorithm, and frankly it’s not obvious to me. Obviously you are pairing some police with some thieves. But it would be helpful if you explained what/how your pairing mechanic was, and why you felt it produced a best-possible result. In cases where there are potential conflicts between pairings, how do you decide which to choose, and why?
Also, note that at the end you are doing a lot of unnecessary work.
new_p = []
new_t = []
for i in p_list:
if i not in new_p:
new_p.append(i)
for j in t_list:
if j not in new_t:
new_t.append(j)
final_combinations = list(zip(new_p, new_t))
print('Number of thieves caught: %s'%(len(final_combinations)))
return len(final_combinations)
First, you make new_p
and new_t
. (Which aren’t new: bad names.) You make them contain unique members of a previous list. Then you zip
them together. The effect of zip
is to stop when the shortest source iterable is exhausted. Then you compute the len
of the result.
Effectively, you’re asking for the length of the shortest unique group of indexes. Since order really doesn’t matter here (you’re using the len
of the result) you can use a set
. And since you just care about the length, you don’t need to zip
them, just compare their lengths:
catches = min(len(set(p_list)), len(set(t_list)))
print('Number of thieves caught: %d' % catches)
return catches