题解 | D和E 题解

from collections import defaultdict,deque
MI = lambda :map(int,input().split())
li = lambda :list(MI())
n,k=li()
a=li()
d=defaultdict(deque)
res=0
l=r=0
while r<n:
    d[a[r]].append(r)
    while len(d[a[r]])>1 and r-d[a[r]][-2]>k:
        d[a[l]].popleft()
        l+=1
    res+=r-l+1
    r+=1
print(res)

from types import GeneratorType
MI = lambda :map(int,input().split())
li = lambda :list(MI())

def bootstrap(f, stack=[]):   #yield
    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

n=li()[0]
g=[[] for _ in range(n)]
for _ in range(n-1):
    u,v=li()
    u-=1;v-=1
    g[u]+=[v]
    g[v]+=[u]

h=[1]*n
res=[0]*n

@bootstrap
def dfs(u,f):
    for v in g[u]:
        if v!=f:
            yield dfs(v,u)
            h[u]=max(h[u],h[v]+1)
    yield None
dfs(0,-1)

@bootstrap
def dfs2(u,f):
    fi=se=0
    for v in g[u]:    #这个时候u是根，所以不需要管f，f也作为u的子节点的
        if h[v]>fi:
            se,fi=fi,h[v]
        elif h[v]>se:
            se=h[v]
    res[u]=fi
    for v in g[u]:
        if v!=f:
            if h[v]==fi:
                h[u]=se+1 #递归前换根，把h[u]换做是v为根时的情况
            else:
                h[u]=fi+1
            yield dfs2(v,u)
    yield None
dfs2(0,-1)

ans=0
for i in res:
    if i<=2:
        ans+=1
print(ans)