ABC-046｜B - AtCoDeerくんとボール色塗り

ボールは直線に並んでいるので最初のボールにはすべての色を選ぶことができますが、それ以降のボールには直前のボールに塗った色以外ということになります。すなわち k (k-1) (k-1) (k-1) ...という並びが全体でn個並ぶことになります。

これを式にしたものが解答なのですが、例題にもあるボールが1個の場合はそのまま色数を返すように場合分けをしています。

 def getIntMap():
    return map(int, input().split())


def main():
    n, k = getIntMap()

    if n == 1:
        print(k)
    else:
        print(k * ((k - 1) ** (n - 1)))


if __name__ == "__main__":
    main()

• 設問ページ