Pythonでn人のスイッチ操作後にオンになっているスイッチの数を求めるプログラム

問題の概要

部屋にn個のスイッチがあり、最初はすべてオフになっています。ここにn人の人が順番に現れ、次のルールでスイッチを切り替えていきます。

• 1番目の人は、1の倍数であるすべてのスイッチ（つまり全スイッチ）を切り替える。
• 2番目の人は、2の倍数であるスイッチ（2、4、6、…）を切り替える。
• i番目の人は、iの倍数であるスイッチを切り替える。

このとき、最終的にオンになっているスイッチの数を求めるのがこの問題の目的です。

具体例：n = 5 の場合

入力が n = 5 のとき、出力は 2 になります。実際の動きを順に確認してみましょう（初期状態はすべてオフ）。

• 初期状態：[0, 0, 0, 0, 0]
• 1番目の人が全スイッチを切り替え：[1, 1, 1, 1, 1]
• 2番目の人が2・4番目を切り替え：[1, 0, 1, 0, 1]
• 3番目の人が3番目を切り替え：[1, 0, 0, 0, 1]
• 4番目の人が4番目を切り替え：[1, 0, 0, 1, 1]
• 5番目の人が5番目を切り替え：[1, 0, 0, 1, 0]

最終的にオンになっているのは1番目と4番目の2つなので、答えは 2 となります。

解法のポイント：平方数に着目する

スイッチ i が最終的にオンになるのは、i の約数の個数が奇数になる場合です。そして、約数の個数が奇数になるのは、i が平方数（1、4、9、16、…）のときだけです。したがって、この問題の答えは「n 以下の平方数の個数」、すなわち ⌊√n⌋ を求めることと同じになります。

アルゴリズム：二分探索で整数平方根を求める

n の整数平方根を二分探索で効率よく求めます。手順は以下の通りです。

• l := 0、r := n と初期化する。
• l ≤ r の間、以下を繰り返す。
  • mid := l + (r − l) / 2 を計算する。
  • mid² ≤ n < (mid + 1)² を満たせば、mid を返す。
  • n < mid² であれば、r := mid とする。
  • それ以外の場合は、l := mid + 1 とする。

Pythonでの実装例

class Solution:
    def solve(self, 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

ob = Solution()
n = 5
print(ob.solve(n))

入力

5

出力

2

計算量と補足

二分探索を用いることで、時間計算量は O(log n)、空間計算量は O(1) と非常に効率的です。なお、Pythonでは標準ライブラリの math.isqrt(n) を使えば、同じ結果を1行で取得することもできます。競技プログラミングや実務においても、この「約数の個数の偶奇」と「平方数」の関係は頻出のテクニックなので、ぜひ覚えておきましょう。