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

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.

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

qstiva.out

0
1

Rezolvare

<syntaxhighlight lang="python"> 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("Invalid input")
       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")

def main() -> None:

   try:
       queries = read_input()
   except ValueError as e:
       print(e, file=sys.stderr)
       sys.exit(1)
   ans = solve_queries(queries)
   write_output(ans)

if __name__ == "__main__":

   sys.exit(main())

</syntaxhighlight>

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".