1924 - QStiva

Cerința

Se dă o stivă inițial vidă. Să se efectueze Q operații de forma:

1 x: Se adaugă x în stivă. 2: Se șterge elementul din vârful stivei. 3 S: Se întreabă dacă se poate scrie valoarea S ca sumă de elemente aflate în stivă. Fiecare element poate fi folosit o singură dată în calcularea sumei. Răspunsul va fi 1 în caz afirmativ și 0 în caz negativ.

Date de intrare

Fișierul de intrare qstiva.in conține pe prima linie numărul Q, iar pe următoarele Q linii se vor afla operațiile descrise mai sus.

Date de ieșire

"Dacă datele sunt introduse corect, pe ecran se va afișa: "Input valid", apoi fișierul de ieșire qstiva.out va conține răspunsurile operațiilor de tipul 3, câte un răspuns pe linie, în ordinea în care acestea apar în fișierul de intrare. În cazul în care datele nu respectă restricțiile, se va doar afișa pe ecran: "Input invalid."

Restricții și precizări

• 1 ⩽ Q ⩽ 100000
• pentru o operație de tipul 1, 1 ≤ x ≤ 1000
• pentru o operație de tipul 3, 1 ≤ S ≤ 1000
• nu se vor efectua operații de tipul 2 sau de tipul 3 dacă stiva este goală.

Exemplu

qstiva.in

8

1 12

2

1 1

3 2

1 2

1 12

3 13

2

qstiva.out

0

1

Ieșire consolă

Input valid

Rezolvare

import sys
from typing import List
from bitarray import bitarray

nmax = 100001
Smax = 1000

def validate_input(queries: List[List[int]]) -> bool:
    if not 1 <= len(queries) <= 100000:
        return False
    stack_size = 0
    for query in queries:
        t = query[0]
        if t == 1:
            x = query[1]
            if not 1 <= x <= 1000:
                return False
            stack_size += 1
        elif t == 2:
            if stack_size == 0:
                return False
            stack_size -= 1
        elif t == 3:
            x = query[1]
            if not 1 <= x <= 1000:
                return False
            if stack_size == 0:
                return False
    return True

def read_input() -> List[List[int]]:
    with open("qstiva.in", "r") as fin:
        Q = int(fin.readline().strip())
        queries = []
        for _ in range(Q):
            query = list(map(int, fin.readline().split()))
            queries.append(query)
        if not validate_input(queries):
            raise ValueError("Input invalid")
        return queries

def solve_queries(queries: List[List[int]]) -> List[int]:
    n = 0
    dp = [bitarray(Smax+2) for _ in range(nmax)]
    dp[0][0] = 1
    ans = []
    for query in queries:
        t = query[0]
        if t == 1:
            x = query[1]
            n += 1
            dp[n].setall(False)
            for i in range(Smax-x+1, Smax+1):
                dp[n][i] = dp[n-1][i]
            for i in range(Smax-x, -1, -1):
                if dp[n-1][i]:
                    dp[n][i+x] = dp[n][i] = True
        elif t == 2:
            n -= 1
        elif t == 3:
            x = query[1]
            ans.append(int(dp[n][x]))
    return ans

def write_output(ans: List[int]) -> None:
    with open("qstiva.out", "w") as fout:
        for a in ans:
            fout.write(str(a) + "\n")

if __name__ == "__main__":
    try:
        queries = read_input()
        print("Input valid", file=sys.stdout)
    except ValueError as e:
        print(e, file=sys.stderr)
        sys.exit(1)
    ans = solve_queries(queries)
    write_output(ans)

Explicație

Funcția validate_input verifică dacă input-ul este valid, adică dacă este în conformitate cu restricțiile problemei. Funcția read_input citește input-ul din fișierul "qstiva.in" și îl validează folosind funcția validate_input. Funcția solve_queries primește input-ul validat și rezolvă problema "Q-Stack" utilizând tehnica de programare dinamică. Funcția write_output scrie rezultatul în fișierul "qstiva.out". Funcția main este funcția principală care apelează celelalte funcții în ordine pentru a citi input-ul, a rezolva problema și a scrie rezultatul.

Algoritmul utilizează o matrice dp de dimensiuni nmax x Smax+2, unde dp[i][j] este True dacă există un subsir de lungime i cu o sumă de elemente egală cu j. Matricea este inițializată cu dp[0][0] = True, pentru că există un subsir de lungime zero cu suma elementelor zero. Algoritmul primește trei tipuri de interogări:

Adaugă elementul x la stack. Elimină ultimul element adăugat la stack. Verifică dacă există un subsir de elemente care au o sumă egală cu x. Pentru fiecare interogare de tipul 1, se creează un nou rând în matricea dp, inițializat cu False. Pentru fiecare element i din matricea anterioară care are valoarea True, se actualizează valorile dp[n][i+x] și dp[n][i] pe baza valorii anterioare dp[n-1][i].

Pentru fiecare interogare de tipul 2, se elimină ultimul rând din matricea dp.

Pentru fiecare interogare de tipul 3, se verifică dacă există un subsir de elemente din stack care au o sumă egală cu x. Dacă da, se adaugă 1 în lista de rezultate ans, altfel se adaugă 0.

În final, lista de rezultate este scrisă în fișierul "qstiva.out".