Submission #65923987

---

Source Code Expand

import sys
MOD = 998244353

def solve():
    input = sys.stdin.read().split()
    ptr = 0
    T = int(input[ptr])
    ptr += 1
    for _ in range(T):
        N = int(input[ptr])
        K = int(input[ptr + 1])
        ptr += 2
        
        s = bin(N)[2:]
        m = len(s)
        # dp[pos][count][tight] = (sum, cnt)
        dp = {}
        
        def dfs(pos, count, tight):
            if pos == m:
                return (0, 1) if count == K else (0, 0)
            key = (pos, count, tight)
            if key in dp:
                return dp[key]
            
            limit = int(s[pos]) if tight else 1
            total_sum = 0
            total_cnt = 0
            for d in range(0, limit + 1):
                new_tight = tight and (d == limit)
                new_count = count + (1 if d == 1 else 0)
                if new_count > K:
                    continue
                sum_sub, cnt_sub = dfs(pos + 1, new_count, new_tight)
                total_sum += sum_sub
                if d == 1:
                    total_sum += cnt_sub * (1 << (m - 1 - pos))
                total_sum %= MOD
                total_cnt += cnt_sub
                total_cnt %= MOD
            dp[key] = (total_sum, total_cnt)
            return dp[key]
        
        res_sum, res_cnt = dfs(0, 0, True)
        print(res_sum % MOD)

solve()

Submission Info

Judge Result