2493 - Recc

Enunț

Ana Mia are o recurență liniară de forma P[N] = A[1]*P[N-1] + A[2]*P[N-2] + A[3]*P[N-3] + A[4]*P[N-4], N ≥ 5. Studiind-o, îi vine o idee MAXIMĂ de problemă: “Pentru câte cvadriplete (P[1], P[2], P[3], P[4]) din mulțimea numerelor naturale [1, B] valoarea P[N] modulo K are valoarea X?”

Cerința

Scrieți un program care citind N, B, X, K, și numerele A[i], 1 ≤ i ≤ 4 rezolvă problema Anei.

Date de intrare

Fișierul de intrare reccIN.txt conține pe prima linie numerele N, B, X, K, separate printr-un spațiu, iar pe a doua linie 4 numere naturale separate prin spații, reprezentând numerele A[i], 1 ≤ i ≤ 4.

Date de ieșire

Fișierul de ieșire reccOUT.txt va conține pe prima linie numărul ANS, răspunsul problemei puse de Ana Mia. În cazul în care restricțiile nu sunt îndeplinite, se va afișa mesajul "Datele nu corespund restrictiilor impuse".

Restricții și precizări

• 5 ≤ N ≤ 10^18
• 1 ≤ B ≤ 1000
• 1 ≤ K ≤ 10^9
• 0 ≤ X < K
• numerele de pe a doua linie a fișierului de intrare vor fi mai mici decât 1.000.000.000

Exemplul 1:

reccIN.txt

6 2 3 4
1 2 3 1

reccOUT.txt

2

Explicație

Soluțiile sunt (2, 1, 2, 1) și (2, 2, 2, 1)

Exemplul 2:

reccIN.txt

6 2 5 4
1 2 3 1

reccOUT.txt

Datele nu corespund restrictiilor impuse

Rezolvare

class Matrice:
    def __init__(self):
        self.M = [[0] * 4 for _ in range(4)]

def inmultire(A, B, K):
    C = Matrice()
    for i in range(4):
        for j in range(4):
            for k in range(4):
                C.M[i][j] = (C.M[i][j] + A.M[i][k] * B.M[k][j]) % K
    return C

def putere(A, N, K):
    if N == 1:
        return A
    if N % 2 == 1:
        return inmultire(A, putere(A, N - 1, K), K)
    P = putere(A, N // 2, K)
    return inmultire(P, P, K)

def verificare_restrictii(N, B, K, X):
    if not (5 <= N <= 10**18):
        return False
    if not (1 <= B <= 1000):
        return False
    if not (1 <= K <= 10**9):
        return False
    if not (0 <= X < K):
        return False
    return True

def F(val, vs, vd, I):
    st, dr = 1, I
    while st <= dr:
        mij = (st + dr) // 2
        if vs[mij] == val:
            return vd[mij]
        if vs[mij] < val:
            st = mij + 1
        else:
            dr = mij - 1
    return 0

def main():
    with open("reccIN.txt", "r") as f:
        # Citim prima linie și împărțim valorile
        linia1 = f.readline().strip().split()
        N = int(linia1[0])
        B = int(linia1[1])
        X = int(linia1[2])
        K = int(linia1[3])
        
        if not verificare_restrictii(N, B, K, X):
            with open("reccOUT.txt", "w") as g:
                g.write("Datele nu corespund restrictiilor impuse\n")
            return
        
        # Citim a doua linie și împărțim valorile
        linia2 = f.readline().strip().split()
        A = [int(val) for val in linia2]
        
    C = Matrice()
    for i in range(4):
        C.M[3 - i][3] = A[i] % K
    
    for i in range(3):
        C.M[i + 1][i] = 1
    
    C = putere(C, N - 4, K)
    
    for i in range(4):
        A[i] = C.M[i][3]
    
    vs = []
    vd = []
    I = 0
    for i in range(1, B + 1):
        for j in range(1, B + 1):
            vs.append((A[0] * i % K + A[1] * j % K) % K)
            I += 1
    
    vs.sort()
    
    vd = [0] * (I + 1)
    vd[0] = 1
    II = 1
    for i in range(1, I):
        if vs[II - 1] != vs[i]:
            vs[II] = vs[i]
            vd[II] = 1
            II += 1
        else:
            vd[II - 1] += 1
    I = II
    
    sol = 0
    for i in range(1, B + 1):
        for j in range(1, B + 1):
            v = (A[2] * i % K + A[3] * j % K) % K
            v = (X - v + K) % K
            sol += F(v, vs, vd, I)
    
    with open("reccOUT.txt", "w") as g:
        g.write(f"{sol}\n")

if __name__ == "__main__":
    main()