HackerRank Algebra – Birthday Triplets

Challenge: Birthday Triplets

Subdomeniu: Algebra (algebra)

Scor cont: 40.0 / 40

Submission status: Accepted

Submission score: 1.0

Submission ID: 464747907

Limbaj: python3

Link challenge: https://www.hackerrank.com/challenges/the-triplets/problem

Cerinta

Julia received a really simple function, $f$, for her birthday! The function is defined as:
$$f_{n} = a^{n} + b^{n} + c^{n}$$ 
Here, $a$, $b$, $c$, and $n$ are positive integers and $a \lt b \lt c$. Unfortunately, she forgot the values of $a$, $b$, and $c$; however, she *does* remember the values of $f_{2}$, $f_{3}$, and $f_{4}$!

Julia wants your help finding the triplet $\left(a,\ b,\ c\right)$ so she can calculate the value of $f_{n}$. If there is more than one such triplet, then she always chooses the one with the smallest value of $a$; if there are still many such triplets, then she chooses the one with the smallest value of $b$.  

----

You are given $q$ queries, where each query consists of $f_{2}$, $f_{3}$, $f_{4}$, $l$, and $r$. For each query, find the value of $S \bmod \left(10^{9} + 7\right)$ and print it on a new line, where $S$ is defined as:   
$$S = \sum_{n = l}^{r}f_{n}$$

**Note:** It is guaranteed that the triplet $\left(a,\ b,\ c\right)$ always exists for the given values of $f_{2}$, $f_{3}$, and $f_{4}$.

Input Format

The first line of the input contains an integer, $q$, denoting the number of queries.	
Each of the $q$ subsequent lines contains five space-separated integers describing the respective values of $f_{2}$, $f_{3}$, $f_{4}$, $l$, and $r$ for a query.

Output Format

For each query, print the value of $S \bmod \left(10^{9} + 7\right)$ on a new line.

Constraints

- $1 \le q \le 2500$
- $6 \le f_{1}\le 15 \times 10^{3}$
- $1 \le l \le r \le 10^{15}$

Cod sursa

#!/bin/python3
import sys
import math

MOD = 10 ** 9 + 7


def recover_triplet(p2: int, p3: int, p4: int):
    best = None

    kmin = math.isqrt(p2)
    if kmin * kmin < p2:
        kmin += 1
    kmax = math.isqrt(3 * p2)

    p2_sq = p2 * p2

    for s1 in range(kmin, kmax + 1):
        if s1**4 - 6 * p2 * s1 * s1 + 8 * p3 * s1 + 3 * p2_sq - 6 * p4 != 0:
            continue

        s2_num = s1 * s1 - p2
        if s2_num & 1:
            continue
        s2 = s2_num // 2

        s3_num = 2 * p3 + s1 * s1 * s1 - 3 * s1 * p2
        if s3_num % 6 != 0:
            continue
        s3 = s3_num // 6

        # Find integer roots a < b < c from:
        # a + b + c = s1, ab + ac + bc = s2, abc = s3.
        for a in range(1, s1 // 3 + 1):
            t = s1 - a              # b + c
            u = s2 - a * t          # bc
            if u <= 0:
                continue

            d = t * t - 4 * u
            if d < 0:
                continue
            sq = math.isqrt(d)
            if sq * sq != d:
                continue
            if (t - sq) & 1:
                continue

            b = (t - sq) // 2
            c = (t + sq) // 2
            if a < b < c and a * b * c == s3:
                cand = (a, b, c)
                if best is None or cand < best:
                    best = cand
                break

    return best


def geom_sum(x: int, l: int, r: int) -> int:
    if x == 1:
        return (r - l + 1) % MOD
    inv = pow(x - 1, MOD - 2, MOD)
    return ((pow(x, r + 1, MOD) - pow(x, l, MOD)) * inv) % MOD


def main():
    data = sys.stdin.buffer.read().split()
    if not data:
        return

    it = iter(data)
    q = int(next(it))
    queries = []
    keys = []
    seen = set()

    for _ in range(q):
        f2 = int(next(it))
        f3 = int(next(it))
        f4 = int(next(it))
        l = int(next(it))
        r = int(next(it))
        queries.append((f2, f3, f4, l, r))
        key = (f2, f3, f4)
        if key not in seen:
            seen.add(key)
            keys.append(key)

    trip = {}
    for key in keys:
        f2, f3, f4 = key
        t = recover_triplet(f2, f3, f4)
        if t is None:
            # Problem guarantees existence, but keep output deterministic.
            t = (1, 2, 3)
        trip[key] = t

    out = []
    for f2, f3, f4, l, r in queries:
        a, b, c = trip[(f2, f3, f4)]
        ans = (geom_sum(a, l, r) + geom_sum(b, l, r) + geom_sum(c, l, r)) % MOD
        out.append(str(ans))

    sys.stdout.write("\n".join(out))


if __name__ == '__main__':
    main()