python 编程 小葱的01串

python L2 研发工程师 编程题 小葱的01串

题目：

给定一个长度为偶数的环形01字符串。（环形指，第一个字符和最后一个字符是相邻的）
字符串初始每个字符都是白色。小葱想把一段连续区间染成红色。使得红色的字符‘0’数量等于白色的字符‘0’数量，红色的字符‘1’数量等于白色的字符‘1’数量。问有多少种不同的染色方法？两个方案不同当且仅当存一个某字符，在一方案是染成红，在另一个方案为白色。
输入描述：
第一行输入一个正整数n,代表字符串长度。
第二行输入一个长度为n的01字符串（仅由字符‘0’和字符‘1’组成的字符串）
数据范围：
2<=n<=30000。保证n是偶数。
输出描述：
合法的染色方案数。

输入
2
11

输出
2
将第一个数字染红为一个方案
将第二个数字染红为一个方案

输入
4
0101
输出：
4

n=int(input())
if (n<2 & n>30000):
    quit()
if (n%2!=0):
    quit()
a1=0
a0=0
a=[0]*(n+5)
b=[0]*(n+5)
s=str(input())
for i in range(n):
    if s[i]==0:
        a0+=1
        a[i+1]=a[i]+1
        b[i+1]=b[i]
    else:
        a1+=1
        a[i+1]=a[i]
        b[i+1]=b[i]+1
for i in range(n):
    if s[i]==0:
        a[i+1+n]=a[i+n]+1
        a[i+1+n]=b[i+n]
    else:
        a[i+1+n]=a[i+n]
        b[i+1+n]=b[i+n]+1
ans=int(n/2)-1
cnt=0
for i in range(1,n+1):
    if (((a[i+ans]-a[i-1])==int(a0/2)) & ((b[i+ans]-b[i-1])==int(a1/2))):
        cnt+=1
print(cnt)