57 lines
1.3 KiB
Python
57 lines
1.3 KiB
Python
#!/bin/python3
|
|
import sys
|
|
from pprint import pprint
|
|
import numpy as np
|
|
import math
|
|
|
|
def manhattan_distance(a, b):
|
|
return np.abs(a - b).sum()
|
|
|
|
def spiral_ccw(A):
|
|
A = np.array(A)
|
|
out = []
|
|
while(A.size):
|
|
out.append(A[0][::-1]) # first row reversed
|
|
A = A[1:][::-1].T # cut off first row and rotate clockwise
|
|
return np.concatenate(out)
|
|
|
|
def base_spiral(nrow, ncol):
|
|
return spiral_ccw(np.arange(nrow*ncol).reshape(nrow,ncol))[::-1]
|
|
|
|
def to_spiral(A):
|
|
A = np.array(A)
|
|
B = np.empty_like(A)
|
|
B.flat[base_spiral(*A.shape)] = A.flat
|
|
return B
|
|
|
|
input_f = sys.argv[1]
|
|
|
|
#########################################
|
|
# #
|
|
# Part 1 #
|
|
# #
|
|
#########################################
|
|
|
|
number = int(input_f)
|
|
org = number
|
|
while not math.sqrt(number).is_integer():
|
|
number+=1
|
|
|
|
mid = int((math.sqrt(number)-1)/2)
|
|
|
|
length = int(math.sqrt(number))
|
|
|
|
a = [mid, mid]
|
|
|
|
arr = to_spiral(np.arange(1,number+1).reshape(length,length))
|
|
|
|
b = np.where(arr == org)
|
|
|
|
print(manhattan_distance(np.array(a),np.array([b[0][0],b[1][0]])))
|
|
|
|
#########################################
|
|
# #
|
|
# Part 2 #
|
|
# #
|
|
#########################################
|