Tối ưu code sau

Tối ưu code sau cho em với ạ, em cảm ơn ạ!
----- Nội dung ảnh -----
```
n=int(input())
A=list(map(int,input().split()))
prefix = [0]*(n+1)
for i in range(1,n+1):
prefix[i]=prefix[i-1]+A[i-1]
dem=0
A=set(prefix)
for i in A:
f=prefix.count(i)
dem+=f*(f-1) / 2
print(int(dem))
```