AtCoder ABC 133 B - Good Distance (Python)

irisruneです。しばらくコンテストは不参加が続きそうですが、今週はB～E問題まで扱う予定です。

問題

atcoder.jp

書いてあることを実装すれば解けますが、B問題にしては実装がやや重いです。

import sys
import math
import numpy as np


def input():
    return sys.stdin.readline().rstrip()


def dist(y, z, d):
    ret = 0.0
    for i in range(d):
        ret += pow(y[i] - z[i], 2)
    return pow(ret, 0.5)


def main():
    n, d = [int(e) for e in input().split()]
    x = [[int(e) for e in input().split()] for _ in range(n)]

    ans = 0
    for i in range(n):
        for j in range(i + 1, n):
            dist_yz = dist(x[i], x[j], d)
            if math.ceil(dist_yz) == math.floor(dist_yz):
                ans += 1
    print(ans)


main()

素直に、各点間の距離（の二乗ノルム）を求めた上で、距離が整数となるかを判定することで解くことができます。距離が整数となるかを判定する方法は、二乗根を求める関数がライブラリにある言語はそれを使えば楽ですし、公式解説のようにfor文で値を探索してもいいです（制約上、距離・ノルムの値がそれほど大きくならないため）。

おそらく実装がこれ以上に軽くなる解き方が存在しないため、二重ループ、ノルムの計算、二乗根（あるいはループによる探索）によりB問題らしからぬ実装の重さとなっています。

雑記

• 最近のB問題が難しくなっているように思いましたが、以前から難しめの問題がたまに出ているようなので気のせいですね、多分。
• 七夕は過ぎましたが、C++から行末のセミコロンを取っ払った言語が欲しいです。