ABC413のD問題 Make Geometric Sequence

ABC413のD問題Make Geometric Sequenceを、解説と他の方の解答を見てなんとか正解できました。等比数列かを調べるより、コーナー・ケースを思いつけずに苦労しました。

D - Make Geometric Sequence

問題： D - Make Geometric Sequence

数列\(A\)が与えられる。これを並べ替えて等比数列になるかを答える問題。

考え方

等比数列は、同じ割合で増減するので、隣との比はどこでも一定。そこで全ての項が、同じ割合で増加または減少することを確認すれば良い。

ただし比を求めるために割り算を使うと誤差が出るので、式を変形して掛け算のみでの関係を調べる。

あとは与えられた数列\(A\)をどの様に並べ替えるかを考える。通常のsort()を使用すると、大きさの順番に並んでしまう。そのため比の値が負の場合、例えば比が-2の\(\{1, -2, 4, -8, 16\}\)をsort()すると、\(\{-8, -2, 1, 4, 16\}\)となってしまう。そこで与えられた数列を絶対値でソートした¹。

比の値が-1の場合

ただし等比数列の比が-1の場合は、絶対値でソートしてもきれいに並びません。例えば、\(\{1, -1, 1, -1\}\)の場合です。

そこで与えられた数列に含まれる値の種類を調べることで対応した。

1種類しかない場合は、全て同じ値なので比が\(1\)の等比数列であることが分かる²。

2種類の場合は、\(\{1, -1, 1, -1\}\)のような数列を考えて、正と負の数が等比数列にできる数か考える。正と負の数の差が2以上、例えば、\(\{1, -1, 1, -1, 1, 1\}\)は等比数列にならない。この時、2種類の値の絶対値が等しいことを調べないと、\(\{1, -2, 1, -2\}\)のようなケースを見逃してしまう。

正と負の数の差は、数列を合計した値から判断する。

• 合計が0の場合は、\(\{1, -1, 1, -1\}\)の様に正の数と負の数が等しい。
• 合計が数列を構成する正の値または負の値と等しい時は、\(\{1, -1, 1\}\)の様にどちらかが1つ足りない場合。
• それ以外の場合は、等比数列ではない。

3種類以上の場合は、絶対値でソートして隣との比が全て等しいか調べる。

これで正解かと思ったのですが、WAが続いて壁にぶち当たりました。

コーナー・ケース

いろいろテストケースを作成してみたのですが、全て正しく判断できます。しかしコードを提出するとほとんどのテストケースがWAになってしまいます。

そこで他の方の正解例を見ていたらchanta1の解答が目に止まり、数列の長さが2の場合の考慮が抜けていることに気がついました。

\(\{1, -2, 1, -2\}\)は等比数列ではありません。しかし\(\{1, -2\}\)は等比数列の一部とみなすことができます。\(\{1, -2, 4\}\)と続くのかも知れません。よって数列の長さが2の場合は、全て等比数列です。

この判定を組み込むことで正解できました。

技術的工夫

List型を絶対値でソートするには、sort(key=abs)またはsorted(List, key=abs)を使用する。

数列に含まれる値が2種類の場合に等比数列を構成できるかは、次の条件を満たす必要が有る。

1. 2種類の値の絶対値が等しい³。
2. 正の値と負の値の数が等しいまたは、正の値と負の値の数の差が1。

条件2は、そのまま正の値と負の値の数を数えても良いが、数列を合計した値を調べることで判断できる。

• 数列の合計が0の場合は、絶対値が等しい正の値と負の値が同数数列に含まれる。よって等比数列の可能性が有る。
• 数列の合計が数列を構成する正の値または負の値と等しい場合は、数列に含まれる絶対値が等しい正の値と負の値の数が1だけ異なる。よって等比数列の可能性が有る。
• 上の2つ以外の場合は、正の値と負の値のどちらかが過剰なので等比数列ではない。

T = int(input())
for _ in range(T):
    N = int(input())
    a = list(map(int, input().split()))

    # Always Yes if the size of array is 2.
    # https://atcoder.jp/contests/abc413/submissions/67411289
    if len(a) == 2:
        print("Yes")
        continue

    a_set = set(a)
    match len(a_set):
        case 1:
            # ratio 1
            print("Yes")
            continue
        case 2:
            # ratio -1
            # count + and -
            delta = sum(a)
            if (sum(a_set) == 0) and ((delta == 0) or (delta in a_set)):
                print("Yes")
            else:
                print("No")
            continue

    b = sorted(a, key=abs)
    # A_i / A_{i-1} = A_{i+1} / A_i
    # (a_i)^2 = A_{i-1} * A_{i+1}
    for i in range(1, len(b) - 1):
        if (b[i] * b[i]) != (b[i - 1] * b[i + 1]):
            print("No")
            break
    else:
        print("Yes")

参考サイト

• D - Make Geometric Sequence 解説
• Submission #67411289 - Denso Create Programming Contest 2025（AtCoder Beginner Contest 413） 数列の長さが2の場合をかんがえる必要を教えてもらいました。