Search code examples
pythonloops

Is there a simpler, more memory-efficient way to solve this 'magnet' problem?


This was my solution to the magnet problem, where you are given a number n and enter n lines of '10' or '01' magnets. Magnets connect with one another on opposite poles and you are tasked with returning a number g, or groups of magnets after they have all connected or not connected.

from itertools import tee

n = int(input())
magnets = []

for i in range(n):
    magnet = input()
    magnets.append(magnet)

def pairwise(iterable):
    a, b = tee(iterable)
    next(b, None)
    return zip(a, b)

m = iter(magnets)

g = 1 # there will always be at least 1 group

for x, nxt in pairwise(m):
    if x != nxt:
        g += 1
    else:
        pass

print(g)

I'm not necessarily concerned with shortening this to the point of it being unreadable, I'm just curious if there is a simpler way to compare each i and i+1 in a simpler way.


Solution

  • Regarding memory: you don't have to store all magnets, just the previous one:

    prev = None
    groups = 0 
    
    n = int(input("count: "))
    for i in range(n):
        magnet = input("either 01 or 10: ")
        if magnet != prev:
            groups += 1
            prev = magnet
    
    print(groups)