結果

問題 No.1597 Matrix Sort
ユーザー qwewe
提出日時 2025-05-14 13:05:11
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 5,545 bytes
コンパイル時間 411 ms
コンパイル使用メモリ 82,780 KB
実行使用メモリ 103,012 KB
最終ジャッジ日時 2025-05-14 13:07:34
合計ジャッジ時間 13,030 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 16 TLE * 11
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
import bisect

# Increase recursion depth just in case, although the solution is iterative.
# sys.setrecursionlimit(200000) 

# If Input/Output speed becomes an issue, uncomment the line below
# input = sys.stdin.readline 

def solve():
    # Read input values N, K, P
    N, K, P = map(int, sys.stdin.readline().split())
    
    # Read array A
    A = list(map(int, sys.stdin.readline().split()))
    
    # Read array B
    B = list(map(int, sys.stdin.readline().split()))

    # Sort array B to enable efficient counting using binary search (bisect module)
    # This takes O(N log N) time.
    B.sort()

    # Define a function `count_le(X)` that counts the number of pairs (i, j)
    # such that the element E[i][j] = (A[i] + B[j]) % P is less than or equal to X.
    # The time complexity of this function is O(N log N) because it iterates through N elements of A
    # and performs two binary searches (bisect operations) on B for each element, each taking O(log N) time.
    def count_le(X):
        # Initialize total count. Python integers handle arbitrary size, so overflow is not an issue for count up to N^2.
        total_count = 0 
        
        # Iterate through each element A[i] in array A
        for i in range(N):
            Ai = A[i]
            
            # Initialize count of valid B[j] for the current A[i]
            current_Ai_count = 0 
            
            # The condition (A[i] + B[j]) mod P <= X is equivalent to checking if A[i] + B[j] falls into
            # one of two possible ranges modulo P:
            # Case 1: 0 <= A[i] + B[j] < P  AND  A[i] + B[j] <= X
            # Case 2: P <= A[i] + B[j] < 2P AND (A[i] + B[j] - P) <= X which simplifies to P <= A[i] + B[j] <= P + X

            # Combining these: we need pairs (i, j) such that:
            # (0 <= A[i] + B[j] <= X) OR (P <= A[i] + B[j] <= P + X)

            # Let's evaluate conditions based on B[j] for a fixed A[i].

            # Evaluate Case 1 equivalent: 0 <= B[j] <= X - A[i]
            # This requires X - A[i] >= 0 for the interval [0, X - A[i]] to be non-empty.
            if X - Ai >= 0:
                # Count number of B[j] in the sorted array B such that 0 <= B[j] <= X - A[i].
                # Since all B[j] >= 0, this is just counting B[j] <= X - A[i].
                # bisect_right(B, val) returns insertion point for val, which is count of elements <= val assuming 0-based indexing.
                idx_upper = bisect.bisect_right(B, X - Ai)
                current_Ai_count += idx_upper

            # Evaluate Case 2 equivalent: P - A[i] <= B[j] <= P + X - A[i]
            # Also need to satisfy the original constraint 0 <= B[j] < P.
            # The intersection gives the range: max(0, P - A[i]) <= B[j] <= min(P - 1, P + X - A[i])
            
            # If Ai == 0, then P - Ai = P. Since B[j] must be < P, this interval is empty for B[j].
            # So Case 2 only contributes if Ai > 0.
            if Ai > 0:
                # Calculate the lower bound for B[j] in this case.
                # Since Ai > 0 and Ai < P, we have 0 < P - Ai < P.
                lower_bound_Bj = P - Ai
                
                # Calculate the upper bound for B[j]. Need B[j] <= P + X - Ai and B[j] <= P - 1.
                upper_bound_Bj = min(P - 1, P + X - Ai)
                
                # Check if the calculated interval [lower_bound_Bj, upper_bound_Bj] is valid (lower <= upper).
                if lower_bound_Bj <= upper_bound_Bj:
                     # Count number of B[j] in the range [lower_bound_Bj, upper_bound_Bj].
                     # This is (count of elements <= upper_bound_Bj) - (count of elements < lower_bound_Bj)
                     # Using bisect module: count = bisect_right(B, upper_bound_Bj) - bisect_left(B, lower_bound_Bj)
                     idx_upper = bisect.bisect_right(B, upper_bound_Bj)
                     idx_lower = bisect.bisect_left(B, lower_bound_Bj)
                     current_Ai_count += idx_upper - idx_lower
            
            # Add the count for the current A[i] to the total count
            total_count += current_Ai_count

        # Return the total count of elements <= X
        return total_count

    # Binary search for the K-th smallest value in the matrix E.
    # The possible values for matrix elements are in the range [0, P-1].
    low = 0
    high = P - 1
    
    # Initialize the answer. P-1 is the maximum possible value an element can take.
    # `ans` will store the smallest `mid` for which `count_le(mid) >= K`.
    ans = P - 1 

    # Perform binary search. The loop runs O(log P) times.
    while low <= high:
        # Calculate midpoint to check. Avoid overflow for large low/high using `low + (high - low) // 2`.
        mid = low + (high - low) // 2
        
        # Calculate how many elements in E are less than or equal to 'mid'.
        num_le_mid = count_le(mid)
        
        # Check if the count is at least K.
        if num_le_mid >= K:
            # If yes, 'mid' is a potential answer. The K-th element could be 'mid' or smaller.
            # Store 'mid' and try searching in the lower half [low, mid-1] for a possibly smaller answer.
            ans = mid
            high = mid - 1 
        else:
            # If count < K, the K-th element must be strictly larger than 'mid'.
            # Search in the upper half [mid+1, high].
            low = mid + 1

    # After binary search, 'ans' holds the K-th smallest value.
    print(ans)

# Execute the solve function
solve()
0