【JavaScript】ABC388まとめ（A~E）

AtCoder Beginner Contest 388

ABCDEの5完（68分9ペナ）でした。
1110perfでレートは962->977（+15）となりました。
D問題までは20分弱で解けていて、E問題も初提出は3分後の23分だったのですが、嘘貪欲にハマってしまったのと、間違ってないのにmultisetの実装が間違っているのではないか？と思い無駄にペナを量産してしまったので、早く解けた割にはパフォーマンスが伸びずに微妙な結果となりました。
今回はA~Eをまとめます。

A - ?UPC

入力の一文字目は英大文字であることが保証されています。
したがって、入力の一文字目と文字列UPCを足して出力しましょう。

function Main(input) {
  input = input.split("\n").map((line) => line.trim());
  const S = input[0];
  console.log(S[0] + "UPC");
}

Main(require("fs").readFileSync(0, "utf8"));

B - Heavy Snake

ヘビは高々$100$匹しかおらず、日数も最大$100$日のため毎日ヘビの重さを更新する全探索が間に合います。
太さ$T$は変わらないので、長さ$L$に経過日数を足したものとで積を取ってヘビの重さの配列$arr$にメモして、1日ごとに$arr$の最大値を取得しましょう。

function Main(input) {
  input = input.split("\n").map((line) => line.trim());
  const [N, D] = input[0].split(" ").map(Number);
  const [T, L] = [[], []];
  for (let i = 0; i < N; i++) [T[i], L[i]] = input[i + 1].split(" ").map(Number);

  // arr[j]:= j番目のヘビの重さ
  const arr = new Array(N).fill(0);
  for (let i = 0; i < D; i++) {
    // i日目のヘビの重さを更新する
    for (let j = 0; j < N; j++) arr[j] = T[j] * (L[j] + i + 1);
    // ヘビの重さの最大を取得する
    console.log(Math.max(...arr));
  }
}

Main(require("fs").readFileSync(0, "utf8"));

C - Various Kagamimochi

$i$番目の餅の小さい方として、二倍以上の大きさの餅とマッチングさせます。
ここで、餅は昇順にソートされているため、二倍以上の大きさの餅以降の餅とは必ずマッチングさせることが可能です。

例

1 2 3 4 5

1と、2以降（2,3,4,5）はマッチング可能
2と、4以降（4,5）はマッチング可能
3,4,5はマッチング不可

この「$idx-1$番目の餅まではマッチングさせることができないが、$idx$番目以降の餅とは必ずマッチングできる」ような$idx$を二分探索で取得すれば、$idx$以降の餅の個数は$N-idx$なので、これがまさに$i$番目の餅を小さい方の餅とした場合の鏡餅の種類数になります。

あとは、これを全ての餅について探索すればよいので、全体としては$O(NlogN)$でこの問題を解くことができました。

function Main(input) {
  input = input.split("\n").map((line) => line.trim());
  const N = Number(input[0]);
  const A = input[1].split(" ").map(Number);

  let ans = 0;
  for (let i = 0; i < N; i++) {
    // i番目の餅について、idx番目以降の餅はペアにすることができる
    let idx = lower_bound(A, 2 * A[i]);
    ans += N - idx;
  }
  
  console.log(ans);
}

const isOK = (arr, mid, key) => arr[mid] >= key;
const lower_bound = (arr, key) => {
  let ng = -1; //「index = 0」が条件を満たすこともあるので、初期値は -1
  let ok = arr.length; // 「index = a.size()-1」が条件を満たさないこともあるので、初期値は a.size()

  while (ok - ng > 1) {
    let mid = Math.floor((ok + ng) / 2);
    if (isOK(arr, mid, key)) ok = mid;
    else ng = mid;
  }
  return ok;
};

Main(require("fs").readFileSync(0, "utf8"));

D - Coming of Age Celebration

この問題を問題文の通り愚直に解いて「成人済みの宇宙人が、新たに成人した宇宙人に一つ石を渡す」処理を行うと$O(N^2)$となり確実にTLEしてしまいます。（遅延セグ木を用いて区間更新をすれば$O(NlogN)$に計算量が抑えられるので間に合いますが、ここでは遅延セグ木を用いた解法は扱いません）

そこで、「一つずつ渡す」という操作ではなく、「持っている石全部か、自分より右側にいる宇宙人の分だけ最終的な石の所持数は減る」という結果に着目します。
そうすれば、$N$人の宇宙人の結果を一回ずつ考えればよいので$O(N)$で済みそうです。

この考え方を採用することにすると、今度は「自分の石を減らした結果だけでなく、右にいる宇宙人の石を増やした結果も伝播させる必要がある」という課題が浮上します。

今回の例だと「自分の一つ右隣」から「自分の持っている石の分だけ右、もしくは右端」までの区間にいる宇宙人に対して石を増やしてあげる必要がありますが、このようにある区間に対して一定の値の増減を効率的に管理する方法としてimos法があります。（imos法の説明はこちらをご覧ください）

具体的な流れを見ていきましょう。
ある宇宙人の持っている石の数は$「最初から持っていた石の個数(A_i)」+「貰った石の個数(nsum)」-「渡した石の個数(pay)」$で表せます。

ここで、$nsum$はimos法で左から順に求めることができ、$pay$は今持っている石の数か、右側にいる宇宙人の人数のどちらか小さい方で表すことができますので、この方法により$O(1)$で各宇宙人の石の個数を確定させることとなります。

したがって、同じ操作を各宇宙人について行っていくため、全体としては$O(N)$で十分高速にこの問題を解くことができました。

function Main(input) {
  input = input.split("\n").map((line) => line.trim());
  const N = Number(input[0]);
  const A = input[1].split(" ").map(Number);

  const sum = new Array(N + 1).fill(0);
  const ans = new Array(N).fill(0);
  
  let nsum = 0;
  for (let i = 0; i < N; i++) {
    // cnt:= 貰った石の個数 + 最初から持っていた石の個数
    // pay:= 渡す石の個数
    // ans[i]:= 最終的に残る石の個数
    let cnt = nsum + A[i];
    let pay = Math.min(cnt, N - i - 1);
    ans[i] = cnt - pay;

    // 次の宇宙人から、pay人後の宇宙人まで石を渡す必要がある
    sum[i + 1]++;
    sum[i + pay + 1]--;

    // 貰った石の個数はimos法で累積的に管理する
    nsum += sum[i + 1];
  }

  console.log(ans.join(" "));
}

Main(require("fs").readFileSync(0, "utf8"));

E - Simultaneous Kagamimochi

この問題はパッと見、小さい順に「二倍以上の餅のうち最小の餅を選ぶ」貪欲でいけそうですが、以下のようなケースで落ちます。

小さい順に「二倍以上の餅のうち最小の餅を選ぶ」の貪欲法
6
1 2 3 4 5 6
→[1,2],[3,6]で2通り×
→実際は[1,4],[2,5],[3,6]の3通り

大きい順に「半分以下の餅のうち最大の餅を選ぶ」貪欲をしても、以下のようなケースで落ちます。

大きい順に「半分以上の餅のうち最大の餅を選ぶ」の貪欲法
6
1 1 3 3 6 6
→[6,3],[6,3]で2通り
→実際は[1,3],[1,6],[3,6]の3通り

ここで何らかの工夫をする必要があるわけですが、鏡餅の小さい方に着目すれば、以下のような事実がわかります。

• 鏡餅の個数は最大で$\lfloor N/2 \rfloor$である
• (効率的に鏡餅を作った時)鏡餅の小さい方となる餅は$\lfloor N/2 \rfloor$番目までの餅に限られる

2番目は少し気づきづらいですが、1番目の制約と、「すでに完成している鏡餅について、小さい方をより小さい餅に代替するか、大きい方をより大きい餅に代替しても鏡餅の条件を満たすこと」に着目すれば自然に発想することができます。

以上で、餅を小さい方のグループと大きい方のグループに分けることができました。
あとは、multisetなどのデータ構造を用いて大きい方のグループの餅を管理して、「大きい方の餅のグループからの削除」と「大きい方の餅のグループからある値以上の最小の餅を取得する」処理を高速にできればこの問題を解くことができます。

以下のACコードでは平衡二分木のmultisetを使っているので、削除をremove（$O(logN)$）、取得をbisectLeft（$O(logN)$）で行っています。
多分ChatGPTにmultisetを作成してと言えば作ってくれるので、興味ある方は加筆して確認してみてください。

function Main(input) {
  input = input.split("\n").map((line) => line.trim());
  const N = Number(input[0]);
  const A = input[1].split(" ").map(Number);
  const B = [...A].filter((_, i) => i >= N / 2);

  // multisetの初期化
  const ms = new MultiSet(B.length, { compress: B });
  for (let i = 0; i < B.length; i++) ms.add(B[i]);

  let ans = 0;
  // N/2以下の餅たちをN/2番目以上の餅とマッチングさせる
  for (let i = 0; i < N / 2; i++) {
    // 半分に分割する場合は、N/2番目以上の最小の餅からマッチングさせる貪欲が通る
    let idx = ms.bisectLeft(2 * A[i]);

    // マッチングできる餅がある場合は採用し、multisetから削除する
    if (ms.get(idx) === -1) continue;
    ans++;
    ms.remove(ms.get(idx));
  }
  
  console.log(ans);
}
// Multisetクラスは割愛