Policemen catch thieves

Posted on

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

Leave a Reply

Your email address will not be published. Required fields are marked *