Pythonでn人がスイッチを切り替えた後に点灯しているライトの数を求めるプログラム

部屋の中にn個のトグルスイッチがあり、そこにn人がいるとします。彼らは以下のルールに従ってスイッチを切り替えていきます。

• 1番目の人は、すべてのスイッチを切り替えます。
• 2番目の人は、2の倍数にあたるスイッチ(2, 4, 6, ...)を切り替えます。
• i番目の人は、iの倍数にあたるスイッチを切り替えます。以降も同様です。

このとき、最終的にONの状態になっているスイッチの数を求める必要があります。

例えば、入力がn = 5の場合、出力は2になります。初期状態では電球は [0, 0, 0, 0, 0] です。

• 1番目の人の操作後: [1, 1, 1, 1, 1]
• 2番目の人の操作後: [1, 0, 1, 0, 1]
• 3番目の人の操作後: [1, 0, 0, 0, 1]
• 4番目の人の操作後: [1, 0, 0, 1, 1]
• 5番目の人の操作後: [1, 0, 0, 1, 0]

最終的にONになっているライトは2つあることがわかります。

解法のアプローチ

この問題を解くために、以下の手順に従います。

• l := 0 とします。
• r := n とします。
• l <= r の間、以下を繰り返します。
  • mid := l + (r - l) / 2 の小数点以下切り捨て
  • mid * mid <= n < (mid + 1)² が成り立つ場合、mid を返します。
  • n < mid² の場合、r := mid とします。
  • それ以外の場合、l := mid + 1 とします。

なぜ平方根で答えが求まるのか?

k番目のスイッチが切り替えられる回数は、kの約数の個数と一致します。そして、約数の個数が奇数になるのは、kが完全平方数の場合だけです。したがって、最終的にONになっているスイッチの数は「n以下の完全平方数の個数」、すなわち ⌊√n⌋ に等しくなります。上記のコードでは、二分探索を用いてnの整数平方根を効率的に計算しています。

実装例

理解を深めるために、以下の実装を見てみましょう。

def solve(n):
    l, r = 0, n
    while l <= r:
        mid = l + (r - l) // 2
        if mid * mid <= n < (mid + 1) * (mid + 1):
            return mid
        elif n < mid * mid:
            r = mid
        else:
            l = mid + 1

n = 5
print(solve(n))

入力

5

出力

2