Search code examples
pythonlogicpython-itertools

In Python, How does the following code work?


INSTRUCTIONS

Prison Break

A prison can be represented as a list of cells. Each cell contains exactly one prisoner. A 1 represents an unlocked cell and a 0 represents a locked cell.

[1, 1, 0, 0, 0, 1, 0]

Starting inside the leftmost cell, you are tasked with seeing how many prisoners you can set free, with a catch. You are the prisoner in the first cell. If the first cell is locked, you cannot free anyone. Each time you free a prisoner, the locked cells become unlocked, and the unlocked cells become locked again.

So, if we use the example above:

[1, 1, 0, 0, 0, 1, 0]
# You free the prisoner in the 1st cell.

[0, 0, 1, 1, 1, 0, 1] 
# You free the prisoner in the 3rd cell (2nd one locked).

[1, 1, 0, 0, 0, 1, 0]
# You free the prisoner in the 6th cell (3rd, 4th and 5th locked).

[0, 0, 1, 1, 1, 0, 1]
# You free the prisoner in the 7th cell - and you are done!
Here, we have set free 4 prisoners in total.

Create a function that, given this unique prison arrangement, returns the number of freed prisoners.

Examples

freed_prisoners([1, 1, 0, 0, 0, 1, 0]) ➞ 4

freed_prisoners([1, 1, 1]) ➞ 1

freed_prisoners([0, 0, 0]) ➞ 0

freed_prisoners([0, 1, 1, 1]) ➞ 0

Notes

  • You are the prisoner in the first cell. You must be freed to free anyone else.

  • You must free a prisoner in order for the locks to switch. So in the second example where the input is [1, 1, 1] after you release the first prisoner, the locks change to [0, 0, 0]. Since all cells are locked, you can release no more prisoners.

  • You always start within the leftmost element in the list (the first prison cell). If all the prison cells to your right are zeroes, you cannot free any more prisoners.

Credits: Edabit

I found the following code in the solutions:

def freed_prisoners(prison):
    freed = [k for k, _ in groupby(prison)]
    return len(freed) if freed[0] == 1 else 0

I know the groupby function will give me keys: [1, 0, 1, 0].

My question is how does this code happen to match the logic? Does it just happen to give the same answer required?


Solution

  • From the prisoner's perspective, the prison break works like this:

    1. If my cell is locked, nobody is freed (freed[0] == 1). The end. Otherwise, continue with step 2.
    2. After I leave my cell, the locks switch, and anybody in my group stays in their cell, since it just went from unlocked to locked. Skip everybody in my group (k, _ in groupby(prison)).
    3. Repeat step 1 and 2 for the rest of the cell block.

    From step 2, we can deduce that only one prisoner gets freed from each group with the same lock status. And we can also deduce that from each following group, we can free one prisoner. So determining the number of prisoners that get freed can be reduced to counting the number of different groups (and zero if my cell is locked). This would be len(freed).

    This can be more easily seen from the code below. Also notice the second example, that addresses the "only one per group freed" part; compare with the first example:

    import itertools
    
    def freed_prisoners(prison):
        groups = [(k,list(v)) for k,v in itertools.groupby(prison)]
        print(f"Groups: {groups}")
        freed = [k for k, _ in groups]
        persons = len(freed) if freed[0] == 1 else 0
        print(f"Freed list: {freed}; freed {persons} persons")
        return persons
    
    freed_prisoners([1, 1, 0, 0, 0, 1, 0])  # ➞ 4
    freed_prisoners([1, 0, 1, 0])           # ➞ 4
    freed_prisoners([1, 1, 1])              # ➞ 1
    freed_prisoners([0, 0, 0])              # ➞ 0
    freed_prisoners([0, 1, 1, 1])           # ➞ 0
    

    Output

    Groups: [(1, [1, 1]), (0, [0, 0, 0]), (1, [1]), (0, [0])]
    Freed list: [1, 0, 1, 0]; freed 4 persons
    Groups: [(1, [1]), (0, [0]), (1, [1]), (0, [0])]
    Freed list: [1, 0, 1, 0]; freed 4 persons
    Groups: [(1, [1, 1, 1])]
    Freed list: [1]; freed 1 persons
    Groups: [(0, [0, 0, 0])]
    Freed list: [0]; freed 0 persons
    Groups: [(0, [0]), (1, [1, 1, 1])]
    Freed list: [0, 1]; freed 0 persons