Find maximum N such that the sum of square of first N natural numbers is not more than X in Python

Suppose we have a given integer X, we have to find the maximum value N so that the sum of squares of first N natural numbers should not exceed the value X.

The sum of squares of first N natural numbers follows the formula: 1² + 2² + 3² + ... + N² = N × (N + 1) × (2N + 1) / 6

So, if the input is like X = 7, then the output will be 2 as 2 is the maximum possible value of N. For N = 3, the sum would be 1² + 2² + 3² = 1 + 4 + 9 = 14, which exceeds X = 7.

Algorithm

We'll use binary search to efficiently find the maximum N:

• Create a function to calculate sum of squares using the formula

• Set search bounds: low = 1, high = large value

• Use binary search to find maximum N where sum ? X

• If current sum ? X, try larger N (move left bound)

• If current sum > X, try smaller N (move right bound)

Implementation

def sum_of_squares(N):
    """Calculate sum of squares of first N natural numbers"""
    return (N * (N + 1) * (2 * N + 1)) // 6

def get_max_N(X):
    """Find maximum N such that sum of squares ? X"""
    low, high = 1, 100000
    result = 0
    
    while low <= high:
        mid = (high + low) // 2
        current_sum = sum_of_squares(mid)
        
        if current_sum <= X:
            result = mid  # Valid N found
            low = mid + 1  # Try larger N
        else:
            high = mid - 1  # Try smaller N
    
    return result

# Test with example
X = 7
max_N = get_max_N(X)
print(f"Maximum N for X = {X}: {max_N}")
print(f"Sum of squares for N = {max_N}: {sum_of_squares(max_N)}")
print(f"Sum of squares for N = {max_N + 1}: {sum_of_squares(max_N + 1)}")

Maximum N for X = 7: 2
Sum of squares for N = 2: 5
Sum of squares for N = 3: 14

How It Works

Let's trace through the example with X = 7:

def trace_binary_search(X):
    """Trace the binary search process"""
    low, high = 1, 10  # Using smaller range for demonstration
    result = 0
    step = 1
    
    print(f"Finding maximum N for X = {X}")
    print("Step | Low | High | Mid | Sum | Action")
    print("-" * 40)
    
    while low <= high:
        mid = (low + high) // 2
        current_sum = (mid * (mid + 1) * (2 * mid + 1)) // 6
        
        if current_sum <= X:
            result = mid
            action = f"Sum ? {X}, try larger"
            low = mid + 1
        else:
            action = f"Sum > {X}, try smaller"
            high = mid - 1
        
        print(f"{step:4} | {low-1 if current_sum <= X else low:3} | {high+1 if current_sum > X else high:4} | {mid:3} | {current_sum:3} | {action}")
        step += 1
    
    return result

trace_binary_search(7)

Finding maximum N for X = 7
Step | Low | High | Mid | Sum | Action
----------------------------------------
   1 |   1 |   10 |   5 |  55 | Sum > 7, try smaller
   2 |   1 |    4 |   2 |   5 | Sum ? 7, try larger
   3 |   3 |    4 |   3 |  14 | Sum > 7, try smaller
   4 |   3 |    2 |   2 |   5 | Sum ? 7, try larger

Testing with Different Values

def test_multiple_values():
    """Test the function with various X values"""
    test_cases = [1, 5, 7, 14, 30, 50]
    
    print("X  | Max N | Sum of Squares")
    print("-" * 30)
    
    for X in test_cases:
        max_N = get_max_N(X)
        sum_val = sum_of_squares(max_N)
        print(f"{X:2} | {max_N:5} | {sum_val:12}")

test_multiple_values()

X  | Max N | Sum of Squares
------------------------------
 1 |     1 |            1
 5 |     2 |            5
 7 |     2 |            5
14 |     3 |           14
30 |     4 |           30
50 |     5 |           55

Conclusion

Binary search efficiently finds the maximum N where the sum of squares doesn't exceed X. The algorithm runs in O(log N) time complexity, making it suitable for large values of X. The key insight is using the mathematical formula N × (N + 1) × (2N + 1) / 6 to calculate the sum directly.