kmjp's blog

競技プログラミング参加記です

LeetCode Weekly Contest 396 : 3139. Minimum Cost to Equalize Array

難易度8の割には簡単?
https://leetcode.com/problems/minimum-cost-to-equalize-array/description/

問題

N要素の整数列Aが与えられる。
以下を行いAの全要素を等しくしたい。
かかる最小コストを求めよ。

  • Aの1要素をcost1のコストでインクリメントする
  • Aのことなる2要素をcost2のコストでインクリメントする

解法

揃える値をmax(A)*2程度まで総当たりする。
極力後者を使おうとした場合のコストをそれぞれ算出していこう。

const ll mo=1000000007;

class Solution {
public:
    int minCostToEqualizeArray(vector<int>& nums, int cost1, int cost2) {
		int N=nums.size();
		sort(ALL(nums));
		ll sum=0;
		int i;
		FOR(i,N) sum+=nums.back()-nums[i];
		ll ret=1LL<<60;
		
		ret=min(ret,sum*cost1);
		for(i=nums.back();i<=2000003;i++) {
			ll a=i-nums[0];
			ll b=sum-a;
			if(a<=b) {
				ret=min(ret,sum/2*cost2+(sum%2)*cost1);
			}
			else {
				ret=min(ret,b*cost2+(a-b)*cost1);
			}
			sum+=N;
		}
		
		
        return ret%1000000007;
    }
};

まとめ

なんかLeetCodeって64bit値で収まる場合も1000000007で割った余りを取らせることあるよな。