class Solution {
    // binary search
    private int bsearch(int x, int[][] pos) {
        int n = pos.length;
        int lptr = 0, rptr = n-1;
        while (lptr <= rptr) {
            int mid = lptr + (rptr - lptr) / 2;
            if (pos[mid][0] == x) return mid;
            if (x < pos[mid][0]) rptr = mid-1;
            else lptr = mid+1;
        }
        return lptr;
    }
    
    public long[] findMaxSum(int[] nums1, int[] nums2, int k) {
        int n = nums1.length;
        
        long[] answer = new long[n];

        // sort indices with position mapping
        int[][] pos = new int[n][2];
        for (int i = 0; i < n; i++) {
            pos[i][0] = nums1[i];
            pos[i][1] = i;
        }
        Arrays.sort(pos, (a, b) -> {
            if (a[0] == b[0]) return a[1] - b[1];
            return a[0] - b[0];
        });
        
        // pref[i] = prefix sum of indices up to i
        long[] pref = new long[n];
        pref[0] = 0 * 1L;

        // maintain k largest so far
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        long sum = 0 * 1L;
        
        for (int i = 1; i < n; i++) {
            int x = nums2[pos[i-1][1]];

            // auto add x
            sum += x;
            pq.add(x);

            // sum at pref[i] = max k previous
            // remove smallest until size = k
            while (pq.size() > k) {
                sum -= pq.poll();
            }
            pref[i] = sum;
        }

        for (int i = 0; i < n; i++) {
            int id = bsearch(nums1[i], pos);
            int x = pos[id][0];
            // find strictly less than nums1[i] on sorted array
            while (id >= 0 && pos[id][0] == x) id--;
            // take prefix sum up to index id
            answer[i] = pref[id+1];
        }
        
        return answer;
    }
}