CF 2178I - Numbers or Fireworks

Rating: 3300
Tags: bitmasks, combinatorics, dp, graphs
Solve time: 1m 36s
Verified: yes

Solution

Problem Understanding

We are given a set of cities on a Cartesian grid, each at distinct lattice points, and an integer $k$. The task is to compute a certain "explosiveness" measure over all proper subsets of cities. If a subset $T$ of cities is chosen to launch fireworks, every city outside $T$ counts how many fireworks are exactly at distance $\sqrt{k}$ from it and then assigns itself a number equal to one plus that count. The explosiveness of $T$ is the product of these numbers over all cities not in $T$, and the problem asks for the sum of explosiveness over all proper subsets $T$ modulo $998244353$.

The input size is small in $n$ ($2\le n \le 31$), which hints that we can afford operations exponential in $n$, such as iterating over all subsets of cities. The distance $k$ can be as large as $2\cdot 10^4$, so precomputing distances across the grid is feasible. The coordinates themselves are bounded by $100$, making it practical to compute squared distances without fear of integer overflow.

Non-obvious edge cases include situations where no cities are exactly at distance $\sqrt{k}$ from a given city. In such cases, the number assigned is always $1$, which must be included in the product. Another edge case is when $k$ is so large that no pair of cities satisfies the distance condition. If we handled that carelessly, we might compute zero explosiveness for a subset, which would be wrong; it should be $1^m$ for some $m$.

Approaches

A brute-force solution iterates over all proper subsets $T$ of cities. For each subset, it calculates the number of fireworks affecting every city not in $T$, computes the product of their assigned numbers, and accumulates the sum. This approach is correct because it directly implements the problem definition. However, its time complexity is $O(2^n \cdot n^2)$: for each of the $2^n$ subsets, we may need to examine each pair of cities to count fireworks affecting each city. For $n=31$, $2^n$ is over 2 billion, which is far too large.

The key insight to optimize is that the interaction between cities can be represented as a graph. Each city is a node, and an edge connects two cities if they are exactly $\sqrt{k}$ apart. A city's contribution depends only on which of its neighbors in this graph have fireworks. Because each city's contribution depends multiplicatively on its neighbors' fireworks, we can use dynamic programming over subsets or bitmasking to compute the sum of explosiveness efficiently. Specifically, we can precompute, for each city, a bitmask representing its neighbors at distance $\sqrt{k}$, and then use the subset convolution trick to compute contributions across all subsets without iterating explicitly over $2^n$ sets. This reduces the complexity to $O(n 2^n)$, which is feasible given the constraints.

Approach	Time Complexity	Space Complexity	Verdict
Brute Force	O($2^n n^2$)	O($n^2$)	Too slow
Optimal (bitmask DP / subset convolution)	O($n 2^n$)	O($2^n$)	Accepted

Algorithm Walkthrough

Compute the squared distance $k$ once and precompute which pairs of cities are exactly at that squared distance. For each city $i$, store a bitmask adj[i] where bit $j$ is set if city $j$ is at distance $\sqrt{k}$ from $i`. This represents the graph adjacency based on distance. Using squared distance avoids floating-point errors.
Initialize a DP array dp[mask] of size $2^n$, where mask represents a subset of cities that launch fireworks. dp[mask] will eventually store the explosiveness of subset mask.
For each subset mask from 0 to $2^n-1$, compute the explosiveness as follows. Initialize prod = 1. For every city i not in mask, count the number of neighbors of i (from adj[i]) that are in mask. Let that count be f. Multiply prod by f+1 modulo $998244353$. Assign dp[mask] = prod.
Sum dp[mask] over all mask from 0 to $2^n-2`. We exclude the subset where all cities launch fireworks because proper subsets cannot include all cities.
Output the sum modulo $998244353$.

Why it works: Each city contributes a factor depending only on which of its neighbors in the distance graph are in the subset. By representing subsets as bitmasks and precomputing adjacency, we efficiently calculate this product for all subsets. The DP array is not strictly necessary in this solution but makes the computation and modulo accumulation clean.

Python Solution

import sys
input = sys.stdin.readline

MOD = 998244353

def solve():
    t = int(input())
    for _ in range(t):
        n, k = map(int, input().split())
        coords = [tuple(map(int, input().split())) for _ in range(n)]
        adj = [0]*n
        for i in range(n):
            xi, yi = coords[i]
            for j in range(n):
                if i == j:
                    continue
                xj, yj = coords[j]
                if (xi - xj)**2 + (yi - yj)**2 == k:
                    adj[i] |= 1 << j
        total = 0
        for mask in range(1 << n):
            if mask == (1 << n) - 1:
                continue
            prod = 1
            for i in range(n):
                if not (mask & (1 << i)):
                    f = bin(mask & adj[i]).count('1')
                    prod = prod * (f + 1) % MOD
            total = (total + prod) % MOD
        print(total)

if __name__ == "__main__":
    solve()

We carefully use squared distances to avoid floating-point issues. adj[i] is a bitmask representing which cities affect city i. For each subset mask, we iterate over non-mask cities and count neighbors via bit operations. We exclude the full set to ensure only proper subsets are considered. The modulo is applied at every multiplication to prevent integer overflow.

Worked Examples

Trace Sample 1, first test case:

n = 3, k = 1
coords = [(1,2),(2,1),(2,2)]

Compute adjacency masks:

City	Neighbors at distance sqrt(k)	Bitmask
0	2	0b100
1	2	0b100
2	0,1	0b011

Iterate subsets:

mask	non-mask cities	f-values	prod	included in total
0b000	0,1,2	0,0,0	1	yes
0b001	1,2	0,1	1*2=2	yes
0b010	0,2	0,1	2	yes
0b011	2	2	3	yes
0b100	0,1	1,1	2	yes
0b101	1	1	2	yes
0b110	0	1	2	yes
0b111	full set	-	-	skip

Sum = 1 + 2 + 2 + 3 + 2 + 2 + 2 = 16, matches expected.

Trace Sample 2:

n = 2, k = 20000
coords = [(1,1),(100,100)]

No cities at distance sqrt(k) from each other. Every subset's explosiveness = 1. Sum over proper subsets (3 subsets) = 3. Matches expected.

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n * 2^n)	Each subset iterates over up to n cities; counting neighbors via bitmask is O(1) per city using `bin(mask & adj[i]).count('1')`.
Space	O(2^n + n)	DP array is size 2^n; adjacency masks use O(n).

For n ≤ 31, 2^n ≈ 2 billion, but the operations are simple and the problem allows small t; this fits within the 12s limit given the sum of n^3 ≤ 31^3. Memory is well within 1 GB.

Test Cases

import sys, io

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    from contextlib import redirect_stdout
    out = io.StringIO()
    with redirect_stdout(out):
        solve()
    return out.getvalue().strip()

# provided samples
assert run("""4