Educatioal Codeforces Round 83-C Adding Powers

内容

• ECF R-83のC問題をpythonで解いた。

  問題

codeforces.com

• 配列 [v_1, .. v_n]が与えられる。各要素は0で初期化されている。
• 配列に対して、下記の操作を行う
  • i-th stepにおいて、(i=0から始まる)
    • 配列の要素を一つ選んで、kⁱを加える　or 何もしない
• 任意の回数操作を行うことができる。
• 与えられた配列を、任意の回数の操作によって作ることができるか?

入力

• T: testcaseの数
• N: 配列の長さ
• K: 加算するときの値。
• A: 作りたい配列

出力

• Aを操作によって作れるならYes　, そうでないならNo

  ポイント

• 逆向きに操作を行う。
  • 「Aから[0,..0]の配列を作れるか」という問題に変わる。
• mをK^mがAの全ての要素よりも大きいという条件を満たす最小の整数とする。
• 繰り返し、以下の操作を行う
  • if (配列の中に、K^mより大きいものがある)
    • 二つ以上あるなら、NO
    • 一つだけなら、そいつからK^mを引いて、continue
    • ０個なら、そのままcontinue
  • m-=1
• 残った配列が全て0ならYES

  実装

from math import ceil, log


def solve():
    N, K = map(int, input().split())
    L = list(map(int, input().split()))

    # K**mが 任意のAの要素より大きくなるようにしたときの最小のm
    m = ceil(log(max(L), K)) if max(L) != 0 else 0
    while m >= 0:
        larger = list(filter(lambda x: x >= K**m, L))
        # 一回の操作で減らせるのは一つまで
        if len(larger) == 1:
            for i, v in enumerate(L):
                if v >= K**m:
                    L[i] -= K**m
                    break
        elif len(larger) >= 2:  # 一回の操作で減らせるのは一つまで
            print("NO")
            return
        m -= 1

    if max(L) != 0:
        print("NO")
    else:
        print("YES")


T = int(input())
for _ in range(T):
    solve()

終わりに

• 逆から辿るのに気付いていれば楽だった。