python - Given A and B, find the number of all integer X values such that X*(X+1) falls between [A, B], both inclusive - Stack O

Given A and B, find the number of all integer X values such that X*(X+1) falls between [A, B], both inclusive.

My approach:

X^2 + X >= A --> X^2 + X - A >= 0

X^2 + X <= B --> X^2 + X - B <= 0

1. Find the roots of both equations using find_quadratic_roots() below.

2. Use the higher and lower roots of both equations to derive all positive & negative X integers (this is where I am stuck)

from math import sqrt, floor, ceil

def find_quadratic_roots(a: int, b: int, c: int) -> tuple:
    """ Returns roots of ax² + bx + c = 0
    For our case, a and b will always be 1 """

    discriminant = b*b - 4*a*c
    if discriminant < 0:
        return None
    
    root1 = (-b + sqrt(discriminant)) / (2*a)
    root2 = (-b - sqrt(discriminant)) / (2*a)

    return (min(root1, root2), max(root1, root2))


def solution(A: int, B: int) -> int:
    if A > B:  return 0
        
    # For lower bound: x² + x - A ≥ 0
    lower_roots = find_quadratic_roots(1, 1, -A)
    
    # For upper bound: x² + x - B ≤ 0
    upper_roots = find_quadratic_roots(1, 1, -B)
    

assert solution(-1, 0) == 2    # -1, 0
assert solution(0, 0) == 2     # -1, 0
assert solution(0, 1) == 2     # -1, 0
assert solution(0, 2) == 4     # -2, -1, 0, 1
assert solution(-5, 5) == 4    # -2, -1, 0, 1
assert solution(0, 6) == 6     # -3, -2, -1, 0, 1, 2
assert solution(6, 6) == 2     # -3, 2
assert solution(3, 6) == 2     # -3, 2

Given A and B, find the number of all integer X values such that X*(X+1) falls between [A, B], both inclusive.

My approach:

X^2 + X >= A --> X^2 + X - A >= 0

X^2 + X <= B --> X^2 + X - B <= 0

1. Find the roots of both equations using find_quadratic_roots() below.

2. Use the higher and lower roots of both equations to derive all positive & negative X integers (this is where I am stuck)

from math import sqrt, floor, ceil

def find_quadratic_roots(a: int, b: int, c: int) -> tuple:
    """ Returns roots of ax² + bx + c = 0
    For our case, a and b will always be 1 """

    discriminant = b*b - 4*a*c
    if discriminant < 0:
        return None
    
    root1 = (-b + sqrt(discriminant)) / (2*a)
    root2 = (-b - sqrt(discriminant)) / (2*a)

    return (min(root1, root2), max(root1, root2))


def solution(A: int, B: int) -> int:
    if A > B:  return 0
        
    # For lower bound: x² + x - A ≥ 0
    lower_roots = find_quadratic_roots(1, 1, -A)
    
    # For upper bound: x² + x - B ≤ 0
    upper_roots = find_quadratic_roots(1, 1, -B)
    

assert solution(-1, 0) == 2    # -1, 0
assert solution(0, 0) == 2     # -1, 0
assert solution(0, 1) == 2     # -1, 0
assert solution(0, 2) == 4     # -2, -1, 0, 1
assert solution(-5, 5) == 4    # -2, -1, 0, 1
assert solution(0, 6) == 6     # -3, -2, -1, 0, 1, 2
assert solution(6, 6) == 2     # -3, 2
assert solution(3, 6) == 2     # -3, 2

• python
• math

Share Improve this question Follow edited Feb 16 at 17:42 bbasaran asked Feb 16 at 16:14 bbasaranbbasaran 40244 silver badges1616 bronze badges 6

• What is the range of A and B? – Paul Hankin Commented Feb 16 at 16:29
• 1 <= A <= B <= 1e9 – bbasaran Commented Feb 16 at 16:51
• Then your first six test cases are all invalid. – no comment Commented Feb 16 at 17:06
• You're right, I missed the constraint until being asked. Even if I ignore the constraint, there should be a solution for negative A and B values as well, since this is math. – bbasaran Commented Feb 16 at 17:08
• I suggest you sketch the curve y=x(x+1) first. Its minimum is -1/4, so there are no values of A or B less than that for which you will find real roots (at which point your assertion is likely to throw a paddy). Then you should confine your assert statement to test cases which are actually valid (though you could lower A and/or B to 0 if you wanted). After that a one-line return statement with judicious uses of floor() and ceil() will solve your problem. – lastchance Commented Feb 17 at 10:26

 |  Show 1 more comment

1 Answer 1

Sorted by: Reset to default 1

You can find all the integer values which satisfy the condition X^2 + X - B <= 0.

Then exclude all the integer values which satify the condition X^2 + X - A < 0.

This is most easily done with sets of integers

def intergers_between_roots(roots: tuple)->range:
    return range(ceil(roots[0]), floor(roots[1]) + 1)

def solution(A: int, B: int) -> int:
    if A > B:  return 0


    # For upper bound only consider values with: x² + x - B ≤ 0
    roots_of_b = find_quadratic_roots(1, 1, -B)
    if roots_of_b is None:
        return 0
    values_below_or_at_B = set(intergers_between_roots(roots_of_b))

    # For lower bound we need to exclude some values
    # Start with: x² + x - A ≤ 0
    roots_of_a = find_quadratic_roots(1, 1, -A)
    if roots_of_a is None:
        return len(values_below_or_at_B)

    values_below_or_at_A = set(intergers_between_roots(roots_of_a))

    # But we need to exclude: x² + x - A < 0
    values_below_A = values_below_or_at_A - set(roots_of_a)


    # We want all the values below or at B but not those below A

    return len(values_below_or_at_B.difference(values_below_A))

Edit: code improvement following comment