From https://www.nicovideo.jp/watch/sm42577409

heat_abnormal.py

import json
import itertools
from typing import Any, Dict, List, Set, Tuple
from tqdm import tqdm
from functools import reduce
import operator
import math


def product_of_positive_factorials(multi_index):
    return reduce(operator.mul, (math.factorial(i) for i in multi_index if i >= 0), 1)


def generate_multi_index_list(
    size: int, degree: int, order: int, limit: int, result_list: List[list[int]] = []
) -> List[List[int]]:
    if size == 1:
        if degree >= order and degree <= limit:
            return [[degree]]
        else:
            return []

    for degree_temp in range(order * (size), degree - order + 1):
        recursive_list: List[List[int]] = generate_multi_index_list(
            size - 1, degree_temp, order, limit, result_list
        )
        recursive_list_correct_size = [
            list for list in recursive_list if len(list) == size - 1
        ]
        result_list_temp = [
            list + [degree - degree_temp]
            for list in recursive_list_correct_size
            if degree - degree_temp <= limit
        ]
        result_list = result_list + result_list_temp
        result_list = sorted(result_list)

    return result_list


def generate_degree_to_multi_indices(
    size: int, degree: int, order: int, limit: int
) -> dict[int, List[List[int]]]:
    degree_range = sorted(range(0, degree + 1), key=abs)

    return {
        lower_degree: generate_multi_index_list(size, lower_degree, order, limit)
        for lower_degree in degree_range
    }


def sort_by_abs_sum(double_list: List[List[int]]) -> List[List[int]]:
    return sorted(double_list, key=lambda x: sum(map(lambda y: abs(y), x)))


class LaurentPoly:
    def __init__(self, size: int, degree: int, order: int, limit: int):
        self.size = size
        self.degree = degree
        self.order = order
        self.degree_to_index_dict = generate_degree_to_multi_indices(
            size, degree, order, limit
        )
        self.multi_indices = sort_by_abs_sum(
            list(itertools.chain.from_iterable(self.degree_to_index_dict.values()))
        )
        key_list = [tuple(list) for list in self.multi_indices]
        self.coeff_dict = {tuple: 0 for tuple in key_list}

    def evaluate(self, variables: list[int]) -> int:
        result = 0
        for exponents, coefficient in self.coeff_dict.items():
            term_value = coefficient
            for i in range(len(exponents)):
                term_value *= variables[i] ** exponents[i]
            term_value = term_value / product_of_positive_factorials(
                self.multi_indices[i]
            )
            result += term_value

        return result

    def evaluate_list(self, variables_list: list[list[int]]) -> list[int]:
        return [self.evaluate(variables_list[i]) for i in range(len(variables_list))]


def get_ranking(value_list: List[int]) -> list[int]:
    ranking = [
        i for i, v in sorted(enumerate(value_list), key=lambda x: x[1], reverse=True)
    ]

    return ranking


def merge_and_count(perm: List[int], start: int, mid: int, end: int) -> int:
    inv_count = 0
    left: list[int] = perm[start : mid + 1]
    right: list[int] = perm[mid + 1 : end + 1]

    i = j = 0
    for k in range(start, end + 1):
        if i < len(left) and j < len(right):
            if left[i] <= right[j]:
                perm[k] = left[i]
                i += 1
            else:
                perm[k] = right[j]
                j += 1
                inv_count += mid - start - i + 1
        elif i == len(left):
            perm[k] = right[j]
            j += 1
        else:
            perm[k] = left[i]
            i += 1
    return inv_count


def count_inversion_recur(perm: List[int], start: int, end: int) -> int:
    if end - start > 0:
        mid = (start + end) // 2
        inv_count = count_inversion_recur(perm, start, mid)
        inv_count += count_inversion_recur(perm, mid + 1, end)
        inv_count += merge_and_count(perm, start, mid, end)

        return inv_count
    else:
        return 0


def count_inversion(perm: List[int]) -> int:
    return count_inversion_recur(perm, 0, len(perm) - 1)


def get_all_negations(numbers: List[int]) -> List[List[int]]:
    negations: Set[Tuple[int, ...]] = set()
    nonzero_indices = [i for i in range(len(numbers)) if numbers[i] != 0]
    for combination in itertools.product([-1, 1], repeat=len(nonzero_indices)):
        negated_list = numbers[:]
        for i in range(len(nonzero_indices)):
            negated_list[nonzero_indices[i]] = (
                numbers[nonzero_indices[i]] * combination[i]
            )
        negations.add(tuple(negated_list))

    return [list(t) for t in negations]


def generate_unit_circle_positive_points(
    size: int, norm: int, result: List[List[int]] = []
) -> List[List[int]]:
    if norm == 0:
        return [[0] * size]

    result_temp = generate_unit_circle_positive_points(size, norm - 1, result)
    result = []
    for inner_list in result_temp:
        for index in range(len(inner_list)):
            new_list = inner_list[:]
            new_list[index] += 1
            if new_list not in result:
                result.append(new_list)

    return result


def generate_unit_circle_all(size: int, norm: int) -> List[List[int]]:
    return sum(
        (
            get_all_negations(list)
            for list in generate_unit_circle_positive_points(size, norm)
        ),
        [],
    )


def distance(point1: List[int], point2: List[int]) -> int:
    if len(point1) != len(point2):
        raise ValueError("Both lists must have the same length.")
    differences = [abs(x - y) for x, y in zip(point1, point2)]

    return sum(differences)


def cancel_neighbors(
    points: List[List[int]], center: List[int], dist: int
) -> List[List[int]]:
    return [point for point in points if distance(point, center) > dist + 1]


def intersection_of_spheres(
    size: int, radius_large: int, center: List[int], small_sphere: List[List[int]]
) -> List[List[int]]:
    shifted_small_sphere = [
        [a + b for a, b in zip(list, center)] for list in small_sphere
    ]
    intersection = [
        lst for lst in shifted_small_sphere if sum(abs(i) for i in lst) == radius_large
    ]

    return intersection


def collate(
    coeffs: List[int],
    laurent_poly: LaurentPoly,
    variables_list: List[Tuple[int, int, int]],
) -> int:
    assign_values([0] + coeffs, laurent_poly.coeff_dict)
    difference = count_inversion(
        get_ranking(
            laurent_poly.evaluate_list(list(map(lambda x: list(x), variables_list)))
        )
    )
    return difference


def coeffs_sorted_by_difference(
    coeffs_list: List[List[int]],
    laurent_poly: LaurentPoly,
    variables_list: List[Tuple[int, int, int]],
) -> List[List[int]]:
    return sorted(coeffs_list, key=lambda x: collate(x, laurent_poly, variables_list))


def extract_top(
    survival_rate: float, survival_limit: int, sorted_coeffs_list: List[List[int]]
) -> List[List[int]]:
    length = len(sorted_coeffs_list)
    number_of_survived = int(length * survival_rate) + 1
    if number_of_survived > survival_limit:
        number_of_survived = survival_limit
    result = sorted_coeffs_list[: number_of_survived + 1]

    return result


def reduce_coeffs_list(coeffs_list: List[List[int]], dist: int) -> List[List[int]]:
    result = []
    reminder = coeffs_list[:]

    while len(reminder) > 0:
        center = reminder[0]
        reminder = cancel_neighbors(reminder, center, dist)
        result.append(center)

    return result


def seach_procedure(  # 原文ママ
    size: int,
    survival_rate: float,
    survival_limit: int,
    radius: int,
    depth: int,
    laurent_poly: LaurentPoly,
    variables_list: List[Tuple[int, int, int]],
) -> Tuple[List[int], int]:
    small_sphere = generate_unit_circle_all(size, radius)

    print("ready...")
    survived_points = [[0] * size]
    for radius_large in tqdm(range(radius)):
        print("searching radius ", radius_large)

        new_survived_points = intersection_of_spheres(
            size, radius_large, survived_points[0], small_sphere
        )
        new_survived_points = coeffs_sorted_by_difference(
            new_survived_points, laurent_poly, variables_list
        )
        new_survived_points = reduce_coeffs_list(new_survived_points, radius)
        new_survived_points = extract_top(
            survival_rate, survival_limit, new_survived_points
        )

        survived_points = new_survived_points[:]

    return (
        survived_points[0],
        collate(survived_points[0], laurent_poly, variables_list),
    )


def extract_values(
    dictionary_list: List[Dict[str, Any]], keys: List[str]
) -> List[Tuple[int, int, int]]:
    result: List[Any] = []
    for dictionary in dictionary_list:
        for key in keys:
            if key not in dictionary.keys():
                return result
        tuple_values = tuple(dictionary[key] for key in keys)
        result.append(tuple_values)
    return result


def assign_values(
    list_values: List[Any], dict_target: Dict[Any, Any]
) -> Dict[Any, Any]:
    for key, value in zip(dict_target.keys(), list_values):
        dict_target[key] = value
    return dict_target


def main(
    degree: int,
    order: int,
    limit: int,
    survival_rate: float,
    survival_limit: int,
    radius: int,
    depth: int,
) -> Any:
    print(
        "degree: " + str(degree),
        "order: " + str(order),
        "limit:" + str(limit),
        "survival rate: " + str(survival_rate),
        "radius: " + str(radius),
        "depth: " + str(depth),
    )

    with open("data.json", "r") as f:
        data_list = json.load(f)
    keys = ["view", "comment", "mylist"]
    variables_list = extract_values(data_list, keys)

    laurent_poly = LaurentPoly(3, degree, order, limit)
    print("laurent_poly is set.")

    size = len(laurent_poly.multi_indices) - 1

    result = seach_procedure(
        size, survival_rate, survival_limit, radius, depth, laurent_poly, variables_list
    )

    print(laurent_poly.multi_indices)

    return result


print(main(4, -1, 10000, 0.05, 200, 4, 12))