Advent of Code 2017, Day 6

aoc_2017_day6.py

INPUT = """14	0	15	12	11	11	3	5	1	6	8	4	9	1	8	4"""

def balance(banks):
	max_blocks = max(banks)
	for idx, v in enumerate(banks):
		if v == max_blocks:
			break

	blocks = banks[idx]
	banks[idx] = 0
	idx += 1
	num_banks = len(banks)
	while blocks:
		if idx == num_banks:
			idx = 0
		banks[idx] += 1
		blocks -= 1
		idx += 1
	return banks
			
def part_one(banks):	
	known_states = set()
	num_rebalances = 0
	tbanks = tuple(banks)
	print num_rebalances, tbanks
	while tbanks not in known_states:
		known_states.add(tbanks)
		balance(banks)	
		tbanks = tuple(banks)
		num_rebalances += 1
		print num_rebalances, banks
	return num_rebalances
	
print part_one([int(v) for v in "0 2 7 0".split()])
print part_one([int(v) for v in INPUT.split()])
		
def part_two(banks):	
	known_states = {}
	num_rebalances = 0
	tbanks = tuple(banks)
	print num_rebalances, tbanks
	while tbanks not in known_states:
		known_states[tbanks] = num_rebalances
		balance(banks)	
		tbanks = tuple(banks)
		num_rebalances += 1
		print num_rebalances, banks
	return num_rebalances - known_states[tbanks]
	
print part_two([int(v) for v in "0 2 7 0".split()])
print part_two([int(v) for v in INPUT.split()])