yukicoder : No.3462 Buttons

場合分けが大変な問題。
https://yukicoder.me/problems/no/3462

問題

整数A,B,Kが与えられる。
変数X=0に対し、以下の処理を順不同で最大K回まで行える場合、Xの最大値を998244353で割った余りを答えよ。

解法

A,Bの正負で丁寧に場合分けする。

• Aが非負の場合
  • Bが2以上の場合、A*B^(K-1)
  • Bが-1,0,1のいずれかの場合、A*K
  • Bが-2以下の場合、Kが奇数ならA*B^(K-1)、偶数なら2A*(B^K-2)
• Aが負の場合、Bも負でKが２以上ならXを正にできる。
  • Bが-1なら(-A)*(K-1)
  • Bが-2以下の場合、Kが偶数ならA*B^(K-1)、奇数なら2A*B^(K-2)

ll T,A,B,K;
const ll mo=998244353;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>A>>B>>K;
		ll ret=0;
		if(A>=0&&B>=0) {
			if(B<=1) ret=A*K%mo;
			else ret=A*modpow(B,K-1)%mo;
		}
		else if(A>=0&&B<0) {
			if(B==-1) {
				ret=A*K%mo;
			}
			else {
				if(K%2==1) ret=A*modpow(-B,K-1)%mo;
				else ret=2*A*modpow(-B,K-2)%mo;
			}
		}
		else if(A<0&&B<0&&K>=2) {
			if(B==-1) {
				ret=(K-1)*(-A)%mo;
			}
			else {
				if(K%2==0) ret=abs(A)*modpow(abs(B),K-1);
				else ret=abs(2*A)*modpow(abs(B),K-2);
			}
		}
		
		
		
		cout<<ret%mo<<endl;
	}
}

まとめ

シンプルな問題ではあるが、場合分けでミスしそう。