B. 弹球 题解

问题描述

有一个长度为 $n$ 的一维网格，第 $i$ 个格子上的字符 $s_i$ 为 '<' 或 '>'。
一颗弹球从某个格子出发，每秒按照当前格子的方向移动一格，移动后该格子上的字符反转（'<' 变 '>'，'>' 变 '<'）。当弹球移出网格（即位置 $<1$ 或 $> n$）时停止。
对于每个起始格子 $i\ (1\le i\le n)$，求弹球离开网格所需的秒数。

输入：多测试用例，总 $n$ 不超过 $5\times 10^5$。
输出：对每个测试用例，一行 $n$ 个整数，表示每个起始位置的答案。

解题思路

直接模拟每个起点会超时（最坏 $O(n^2)$）。注意到弹球的移动具有“往返”性质，每个格子被经过时方向会翻转，这等价于弹球在左右两侧的“墙”之间来回反射，而墙的位置由初始的 '<' 和 '>' 的分布决定。

关键转化

设当前位置为 $i$，方向为 '>'。弹球向右移动一格后，原格子变成 '<'，相当于弹球“穿过”了这个格子并把它变成相反方向。从整体上看，一个 '>' 相当于“向右的力”，'<' 相当于“向左的力”。弹球运动的过程可以看作：它总是朝着当前格子的方向走，但经过的格子会“反转”，所以后续经过时会改变方向。

另一种经典模型：把格子上的方向看作箭头，弹球移动时箭头掉头。可以等价地认为弹球在左右边界之间来回运动，每遇到一个箭头就反弹一次，而箭头方向的变化恰好使得每次反弹的边界向中间收缩。

数学推导

定义：

• $cntL(i)$ 表示 $[1,i]$ 中 '<' 的个数
• $cntR(i)$ 表示 $[1,i]$ 中 '>' 的个数
• $sumL(i)$ 表示 $[1,i]$ 中所有 '<' 的位置之和
• $sumR(i)$ 表示 $[1,i]$ 中所有 '>' 的位置之和

对于起点 $i$，分两种情况：

情况1：$s_i = \text{'>'}$（向右出发）

设右侧 '<' 的个数为 $R_< = cntL(n) - cntL(i)$，左侧 '>' 的个数为 $L_> = cntR(i-1)$。
比较 $L_>$ 和 $R_<$ 的大小：

• 若 $L_> > R_<$：弹球最终会从 右侧 离开网格。
  它向右走的过程中，会先经历若干次来回，最终弹射出去。计算可得答案为

  $$2\Bigl((sumL(n)-sumL(i)) - (sumR(i)-sumR(p-1))\Bigr) + i + (n+1) $$

  其中 $p$ 是左边第 $L_> - R_<$ 个 '>' 的位置（通过二分查找得到）。

• 否则：弹球从 左侧 离开。答案为

  $$2\Bigl((sumL(p)-sumL(i)) - (sumR(i)-sumR(0))\Bigr) + i $$

  其中 $p$ 是右边第 $R_< - L_> + 1$ 个 '<' 的位置。

情况2：$s_i = \text{'<'}$（向左出发）

对称处理。设左侧 '>' 的个数为 $L_> = cntR(i-1)$，右侧 '<' 的个数为 $R_< = cntL(n)-cntL(i)$。
比较 $L_>$ 和 $R_<$ 的大小：

• 若 $L_> \ge R_<$：从 右侧 离开。
  答案为

  $$2\Bigl((sumL(n)-sumL(i-1)) - (sumR(i)-sumR(p-1))\Bigr) - i + (n+1) $$

  其中 $p$ 为左边第 $L_> - R_< + 1$ 个 '>' 的位置。

• 否则：从 左侧 离开。
  答案为

  $$2\Bigl((sumL(p)-sumL(i-1)) - (sumR(i)-sumR(0))\Bigr) - i $$

  其中 $p$ 为右边第 $R_< - L_>$ 个 '<' 的位置。

所有公式中的 sumR(0) 和 p-1 等边界情况需要特殊处理，代码中通过预处理前缀和以及二分查找实现。

复杂度

预处理 $O(n)$，每个询问二分查找 $O(\log n)$，总复杂度 $O(n\log n)$，可以通过 $5\times10^5$ 的数据。

参考代码（标程）

#include <map>
#include <set>
#include <cmath>
#include <ctime>
#include <queue>
#include <stack>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <cstring>
#include <algorithm>
#include <iostream>
using namespace std;
typedef double db; 
typedef long long ll;
typedef unsigned long long ull;
const int N=1000010;
const int LOGN=28;
const ll  TMD=0;
const ll  INF=2147483647;
int  T,n;
ll   Sl[N],Sr[N],IDl[N],IDr[N];
char s[N];

int findpre(int x)
{
	int L=0,R=n+1,M;
	while(L+1!=R)
	{
		M=(L+R)>>1;
		if(Sr[M]<x) L=M;
		else R=M; 
	}
	return R;
}

int findsuf(int x)
{
	int L=0,R=n+1,M;
	while(L+1!=R)
	{
		M=(L+R)>>1;
		if(Sl[n]-Sl[M-1]<x) R=M;
		else L=M; 
	}
	return L;
}

int main()
{
	scanf("%d",&T);
	while(T--)
	{
    	        scanf("%d%s",&n,s);
		for(int i=1;i<=n;i++) 
		{
	  	  Sr[i]=Sr[i-1]+(s[i-1]=='>');
	  	  Sl[i]=Sl[i-1]+(s[i-1]=='<');
 		  IDr[i]=IDr[i-1]+i*(s[i-1]=='>'); 	
 		  IDl[i]=IDl[i-1]+i*(s[i-1]=='<'); 	
		}		
		for(int i=1;i<=n;i++)
		{
			if(s[i-1]=='>')
			{
				if(Sr[i]>Sl[n]-Sl[i])
				{
					int p=findpre(Sr[i]-(Sl[n]-Sl[i]));
					printf("%I64d ",2*((IDl[n]-IDl[i])-(IDr[i]-IDr[p-1]))+i+(n+1));
				}
				else
				{
					int p=findsuf((Sl[n]-Sl[i])-Sr[i]+1);
					printf("%I64d ",2*((IDl[p]-IDl[i])-(IDr[i]-IDr[0]))+i);
				}
			}
			else
			{
				if(Sr[i]>=Sl[n]-Sl[i-1])
				{
					int p=findpre(Sr[i]-(Sl[n]-Sl[i-1])+1);
					printf("%I64d ",2*((IDl[n]-IDl[i-1])-(IDr[i]-IDr[p-1]))-i+(n+1));
				}
				else
				{
					int p=findsuf((Sl[n]-Sl[i-1])-Sr[i]);
					printf("%I64d ",2*((IDl[p]-IDl[i-1])-(IDr[i]-IDr[0]))-i);
				}
			}
		}
		printf("\n");
	}
	
	return 0;
}

注意事项

• 数组大小开至 $N=1000010$ 是为了处理两个方向的前缀和以及二分查找中的边界。
• 输出使用 %I64d（适用于 Windows 下的 long long），若评测环境为 Linux 可改为 %lld。
• 二分查找函数 findpre 和 findsuf 分别用于找到第 $x$ 个 '>' 的位置和第 $x$ 个 '<' 的位置（从右往左数）。