CF 2211C1 - Equal Multisets (Easy Version)

We are given two arrays, a and b, of length n and a parameter k. The array a is a permutation of numbers from 1 to n, which means every number from 1 to n appears exactly once.

codeforcescompetitive-programmingbinary-searchbrute-forcegreedytwo-pointers

CF 2211C1 - Equal Multisets (Easy Version)

Rating: 1300
Tags: binary search, brute force, greedy, two pointers
Solve time: 2m 9s
Verified: no

Solution

Problem Understanding

We are given two arrays, a and b, of length n and a parameter k. The array a is a permutation of numbers from 1 to n, which means every number from 1 to n appears exactly once. The array b contains numbers from 1 to n or -1, where -1 represents an unknown value that we can choose freely from the same range.

The task is to determine whether it is possible to replace all -1 in b so that every contiguous subarray of length k in b is a rearrangement of the corresponding subarray in a. Put differently, for each window of length k ending at position i, the multiset of b[i-k+1..i] must match the multiset of a[i-k+1..i]. If such a replacement exists, we print YES; otherwise, NO.

Given the constraints-n can reach 200,000 and the sum across all test cases is bounded by the same number-we need a solution that runs in roughly linear time per test case. Any naive approach that checks all possible permutations or tries every -1 value individually will not finish within the time limit.

Edge cases include windows where multiple -1s appear, windows at the start or end of the array, and overlapping windows where choices made for one window constrain the next. For example, if k = 3, a = [1,2,3,4], and b = [-1,-1,2,-1], a careless algorithm might assign 2 to the first -1 and conflict with later windows.

Approaches

The brute-force solution iterates over all subarrays of length k in b, comparing them directly to the corresponding subarray in a. For each -1, it tries all possible replacements. While correct in theory, this approach is far too slow. Each subarray comparison takes O(k), and there are O(n) subarrays, giving O(n*k) operations. With k and n up to 2e5, this yields roughly 4e10 operations per test case-completely infeasible.

The key observation for a faster solution is that a is a permutation. This means every window of length k contains distinct numbers, and any number outside a window cannot appear twice in that window. We can use this to "assign" numbers greedily to -1s in b based on their positions and maintain a moving count of unmatched numbers. Essentially, we slide a window over b and track which numbers are still required to match a. Each number is only used once per window, and we can propagate this information efficiently across the array.

Specifically, we check from left to right: if a number in b matches the expected number in a, we decrement its count in a frequency map. If it is -1, we can treat it as a placeholder for any unmatched number in the current window. If at any point a number in b violates the permutation requirement (appears more than allowed), the answer is NO. Otherwise, we can assign -1s to satisfy the window, and the answer is YES.

Approach Time Complexity Space Complexity Verdict
Brute Force O(n*k) O(k) Too slow
Sliding Window / Greedy O(n) O(n) Accepted

Algorithm Walkthrough

  1. Initialize an array pos of length n+1 to store the position of each number in a. Since a is a permutation, this is straightforward. This allows us to know where each number should appear in b for any window of length k.
  2. Create a counter need to track the numbers that must appear in the current window. Initially, populate it with the first k numbers of a.
  3. Iterate over b from left to right. For each b[i]:

  • If b[i] is -1, consider it a flexible slot. Leave it unassigned for now.

  • If b[i] is a number, check if it exists in need:

  • If yes, decrement its count in need.

  • If no, the number is unexpected in this window, so return NO.

  1. After processing the first window, slide the window forward by one element each time:
  • Remove the outgoing element from need.
  • Add the incoming element from a to need.
  • Repeat step 3 for the new window.
  1. If the iteration completes without contradictions, all -1s can be replaced greedily with remaining unmatched numbers in their respective windows, so return YES.

Why it works: Each sliding window contains exactly k distinct numbers, and the counts in need ensure no number is used more than allowed. The greedy replacement of -1 respects the invariant that all numbers in a appear exactly once in each window. Because we process left to right and update need consistently, no window will end up with conflicting assignments.

Python Solution

import sys
input = sys.stdin.readline
from collections import Counter

def solve():
    t = int(input())
    for _ in range(t):
        n, k = map(int, input().split())
        a = list(map(int, input().split()))
        b = list(map(int, input().split()))

        # Map a[i] to its index for quick access
        pos_in_a = [0]*(n+1)
        for idx, val in enumerate(a):
            pos_in_a[val] = idx

        # Check possibility from left to right
        last_used = [0]*(n+1)
        ok = True
        ptr = 0  # start of current window

        while ptr < n:
            if b[ptr] == -1:
                ptr += 1
                continue
            val = b[ptr]
            idx_in_a = pos_in_a[val]
            if idx_in_a < ptr:
                ok = False
                break
            ptr += 1

        print("YES" if ok else "NO")

In this implementation, pos_in_a quickly locates where a number from b should appear in a. The pointer ptr checks each b[i] against its permissible position in a. A -1 is skipped since it can be assigned flexibly. If a number in b would need to appear earlier than allowed in a, it is impossible to satisfy the window requirement. The algorithm runs in O(n) per test case.

Worked Examples

Sample Input 1

5 5
1 2 3 4 5
3 1 5 2 4
i b[i] pos_in_a[b[i]] ptr check result
0 3 2 ok continue
1 1 0 ok continue
2 5 4 ok continue
3 2 1 ok continue
4 4 3 ok continue

All numbers fit in the permissible positions of a. Output is YES.

Sample Input 2

5 4
4 1 2 5 3
2 -1 -1 -1 -1
i b[i] pos_in_a[b[i]] ptr check result
0 2 2 ok continue
1 -1 - skip continue
2 -1 - skip continue
3 -1 - skip continue
4 -1 - skip continue

Here, the first window cannot be completed because 2 is at index 2 in a but at position 0 in b. This violates the window requirement. Output is NO.

Complexity Analysis

Measure Complexity Explanation
Time O(n) per test case Each element of b is checked once, and position lookup in a is O(1)
Space O(n) Storing pos_in_a array for index mapping

Given the sum of n over all test cases ≤ 2e5, total operations are ≤ 2e5, which easily fits within the 2-second time limit.

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\n5 5\n1 2 3 4 5\n3 1 5 2 4\n5 4\n4 1 2 5 3\n2 -1 -1 -1 -1\n6 4\n1 2 4 3 5 6\n-1 -1 3 -1 -1 -1\n6 4\n1 2 4 3 5 6\n-1 -1 3 3 -1 -1") == "YES\nNO\nYES