[PROJECT_EULER] 61. Cyclical figurate numbers

   문제는 이곳에서 확인할 수 있다.

   계산을 해보니 4자리의 polygonal number는 각각 100여개가 채 되지 않았다. 또한 3부터 8까지의 숫자를 나열하는 개수도 많지 않다. 이 정도면 brute force로도 해결이 가능할 것 같다.

import pandas as pd
import itertools 
import functools

def p(k, n):
    
    if k == 3:
        return n*(n+1)//2
    
    if k == 4:
        return n*n
    
    if k == 5:
        return n*(3*n-1)//2
    
    if k == 6:
        return n*(2*n-1)
    
    if k == 7:
        return n*(5*n-3)//2
    
    if k == 8:
        return n*(3*n-2)

p_num = dict()

for k in range(3, 9):
    n = 1
    _p_nums = []
    
    while True:   
        _str = str(p(k, n))
        _len = len(_str)
        
        if _len == 4 and _str[2] != '0':  # 3번째 수가 0이 되면 cycle을 만들 수 없음 
            _p_nums.append([_str[:2], _str[2:]])
        
        if _len >= 5:
            p_num[k] = pd.DataFrame(_p_nums, columns=[f'head{k}', f'tail{k}'])
            break
        
        n += 1
        
for k_cand in itertools.permutations(range(3, 9)):
    
    res = functools.reduce(lambda x, y: x.merge(y, 
                                                left_on=x.columns[-1], 
                                                right_on=y.columns[0]),
                           (p_num[k] for k in k_cand))
    
    head, tail = res.columns[0], res.columns[-1]
    common = set(res[head]) & set(res[tail])
    
    if common != set():
        res = res.query(f'{head} in @common and {tail} in @common')
        ans = res.apply(lambda x: int(x) if 'tail' in x.name else int(x)*100)
        ans = sum(ans)
        
        print(res)
        print(ans)

        break

   6개의 순서쌍은 {8256, 5625, 2512, 1281, 8128, 2882}이며, 답은 28684이다.