#1158435 (Python3) No.3506 All Distance is Square Number

• 問題: No.3506 All Distance is Square Number
• ユーザー: sient
• 提出日時: 2026-04-17 23:28:34
• 言語: Python3 (3.14.3 + numpy 2.4.4 + scipy 1.17.1)
• 結果: AC
• 実行時間: 215 ms / 2,000 ms
• コード長: 5,449 bytes
• コンパイル時間: 568 ms
• コンパイル使用メモリ: 20,956 KB
• 実行使用メモリ: 25,204 KB
• 最終ジャッジ日時: 2026-04-17 23:29:19
• 合計ジャッジ時間: 6,546 ms
• ジャッジサーバーID （参考情報）: judge3_0 / judge1_0

ファイルパターン	結果
other	AC * 29

ソースコード

#!/usr/bin/env python3
import random
import sys
from collections import defaultdict
from math import isqrt


SPECIAL_CYCLE_CASES = {
    5: {
        "bundles": [1, 1, 1, 1, 3],
        "weights": [7, 6, 1, 5, 4, 2, 3],
    },
    6: {
        "bundles": [1, 1, 2, 1, 2, 2],
        "weights": [8, 6, 4, 3, 5, 2, 1, 7, 9],
    },
    7: {
        "bundles": [1, 2, 2, 2, 2, 1, 1],
        "weights": [1, 9, 2, 4, 7, 3, 8, 11, 10, 5, 6],
    },
    8: {
        "bundles": [2, 1, 2, 2, 2, 2, 1, 1],
        "weights": [12, 13, 2, 8, 4, 3, 11, 1, 7, 10, 6, 5, 9],
    },
    9: {
        "bundles": [2, 1, 2, 1, 2, 2, 1, 2, 2],
        "weights": [3, 10, 13, 4, 12, 7, 6, 14, 11, 1, 15, 2, 5, 9, 8],
    },
    10: {
        "bundles": [1, 2, 2, 2, 2, 2, 2, 1, 1, 2],
        "weights": [10, 2, 8, 1, 4, 3, 12, 15, 13, 6, 7, 14, 11, 17, 5, 16, 9],
    },
    11: {
        "bundles": [2, 1, 1, 2, 2, 2, 2, 2, 2, 2, 1],
        "weights": [8, 12, 15, 18, 9, 14, 4, 6, 1, 13, 19, 11, 16, 7, 10, 3, 2, 5, 17],
    },
}


def square_mask(limit: int) -> int:
    mask = 0
    for x in range(1, isqrt(limit) + 1):
        mask |= 1 << (x * x)
    return mask


def build_manual_case(n: int):
    if n == 2:
        edges = [(1, 2, 1)]
        paths = {(1, 2): [1]}
        return edges, paths

    if n == 3:
        edges = [(1, 2, 16), (2, 3, 9), (1, 3, 25)]
        paths = {
            (1, 2): [1],
            (1, 3): [3],
            (2, 3): [2],
        }
        return edges, paths

    if n == 4:
        edges = [
            (1, 3, 18),
            (1, 4, 27),
            (2, 3, 4),
            (1, 2, 36),
            (3, 4, 9),
        ]
        paths = {
            (1, 2): [4],
            (1, 3): [2, 5],
            (1, 4): [4, 3, 5],
            (2, 3): [3],
            (2, 4): [3, 1, 2],
            (3, 4): [5],
        }
        return edges, paths

    raise ValueError("manual case only supports N <= 4")


def build_cycle_case(n: int):
    if n in SPECIAL_CYCLE_CASES:
        bundles = SPECIAL_CYCLE_CASES[n]["bundles"]
        weights = SPECIAL_CYCLE_CASES[n]["weights"]
    else:
        bundles = [1, 1, 1] + [2] * (n - 3)
        weights = list(range(1, 2 * n - 2))
        random.Random(n * 1000).shuffle(weights)

    edges = []
    segment_edges = [[] for _ in range(n + 1)]
    idx = 0
    for seg in range(1, n + 1):
        u = seg
        v = 1 if seg == n else seg + 1
        for _ in range(bundles[seg - 1]):
            w = weights[idx]
            idx += 1
            edge_id = len(edges) + 1
            edges.append((u, v, w))
            segment_edges[seg].append((edge_id, w))

    return edges, segment_edges


def advance(v: int, steps: int, n: int) -> int:
    return ((v - 1 + steps) % n) + 1


def build_reachability(n: int, segment_edges):
    max_sum = 200 * n
    sq = square_mask(max_sum)
    reach = [[0] * (n + 1) for _ in range(n + 1)]
    reach_sq = [[False] * (n + 1) for _ in range(n + 1)]

    for start in range(1, n + 1):
        bits = 1
        reach[start][start] = 1
        for steps in range(1, n):
            seg = advance(start, steps - 1, n)
            new_bits = 0
            for _, w in segment_edges[seg]:
                new_bits |= bits << w
            bits = new_bits
            end = advance(start, steps, n)
            reach[start][end] = bits
            reach_sq[start][end] = (bits & sq) != 0

    return reach, sq


def recover_arc(start: int, end: int, target: int, n: int, segment_edges, reach):
    path = []
    cur = start
    while cur != end:
        seg = cur
        nxt = 1 if cur == n else cur + 1
        for edge_id, w in segment_edges[seg]:
            if target < w:
                continue
            if ((reach[nxt][end] >> (target - w)) & 1) == 0:
                continue
            path.append(edge_id)
            target -= w
            cur = nxt
            break
        else:
            raise RuntimeError("failed to reconstruct cycle path")
    return path


def build_cycle_paths(n: int, segment_edges):
    reach, sq = build_reachability(n, segment_edges)
    paths = {}

    for i in range(1, n):
        for j in range(i + 1, n + 1):
            forward_bits = reach[i][j]
            backward_bits = reach[j][i]
            picked = None

            hit = forward_bits & sq
            if hit:
                target = hit.bit_length() - 1
                picked = recover_arc(i, j, target, n, segment_edges, reach)
            else:
                hit = backward_bits & sq
                if hit:
                    target = hit.bit_length() - 1
                    picked = list(reversed(recover_arc(j, i, target, n, segment_edges, reach)))

            if picked is None:
                raise RuntimeError(f"no square path for pair {(i, j)}")
            paths[(i, j)] = picked

    return paths


def solve(n: int):
    if n <= 4:
        return build_manual_case(n)

    edges, segment_edges = build_cycle_case(n)
    paths = build_cycle_paths(n, segment_edges)
    return edges, paths


def main():
    n = int(sys.stdin.readline())
    edges, paths = solve(n)

    out = [str(len(edges))]
    for u, v, w in edges:
        out.append(f"{u} {v} {w}")
    for i in range(1, n):
        for j in range(i + 1, n + 1):
            path = paths[(i, j)]
            out.append(str(len(path)) + (" " + " ".join(map(str, path)) if path else ""))

    sys.stdout.write("\n".join(out))


if __name__ == "__main__":
    main()