CodeIQ:芋づる式退職

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「芋づる式退職」(CodeIQ)を参照してください。

問題の概要

問題を引用します。
開発者には一人だけ尊敬する同僚がいるものとします。
尊敬する同僚が退職すると自身も退職します。
この条件を前提とし、ある社員1名がやめた場合に残る社員数を出力するプログラムを作成してください。

【ルール】
  • 標準入力は10行です(10行目に改行がつくため、末尾の空行も含めると11行)
  • 標準入力の1行目に退職者の名前が[a-z]からなる3文字の文字列で与えられます
  • 標準入力の2-10行目に社員の名前と尊敬する社員の名前がともに[a-z]からなる3文字の文字列でカンマ区切りで与えられます
  • 1行目の退職者の名前は2-10行目の社員の名前(カンマ区切りの1項目目)には含まれません
  • 退職者および社員の名前は一意です
  • 標準入力の末尾には改行があります
  • 標準出力は会社に残った社員の人数を出力してください
  • 標準出力の末尾に改行をつけてください

私のプログラム

Javaで解答しています。

import java.io.*;
import java.util.*;

public class Main{
	public static class Member{
		Map<String, Set<String>> Mem;

		public Member(){
			Mem = new HashMap<>();
		}

		public void put(String boss, String name){
			if(Mem.containsKey(boss)){
				Set<String> s = Mem.get(boss);
				s.add(name);
			}
			else{
				Set<String> s = new HashSet<>();
				s.add(name);
				Mem.put(boss, s);
			}
		}

		public int remove(String boss){
			Deque<String> q = new LinkedList<>();
			q.addLast(boss);

			while(!q.isEmpty()){
				String b = q.remove();
				Set<String> s = Mem.remove(b);
				if(s != null){
					q.addAll(s);
				}
			}
			return getCount();
		}

		public int getCount(){
			int c = 0;
			for(Map.Entry<String, Set<String>> e: Mem.entrySet()){
				c += e.getValue().size();
			}
			return c;
		}
	}

	public static void main(String[] args) throws IOException{
		String fp = "";	// 最初の退職者
		Member mem = new Member();

		try(BufferedReader br = new BufferedReader(new InputStreamReader(System.in))){
			String line;

			for(int i=0; (line = br.readLine()) != null; i++){
				line = line.trim();
				if(i==0){
					fp = line;
				}
				else{
					String[] p = line.split(",");
					mem.put(p[1], p[0]);
				}
			}
			System.out.println(mem.remove(fp));
		}
	}
}

解説

簡単な問題です。コツはデータをどのように管理するかです。
このプログラムの本体はMemberクラスなのでMemberクラスを中心に説明します。
※問題が途中で補足され若干修正されているようです。私は修正前に問題を解いたので修正が影響しているかもしれません。

Memメンバ変数

入力値を保持する変数です。Map<String, Set<String>>と少しややこしい格好をしていますが、C++のMultiMapのような扱いになります。キーが尊敬される同僚の名前で値のSetにはその同僚を慕う社員の名前が入ります。
つまり、入力値は「社員, 尊敬する同僚」ですが、Memは「尊敬される同僚, [同僚を尊敬する社員, …]」と言うデータを保持するというわけです。

put()

Memに入力値をセットします。キーがすでに存在したら同一キーの値に社員の名前を追加します。

remove()

1行めに入力された名前を削除し、それを慕う社員を削除する、という処理を連鎖的に行います。データ構造で工夫しているので処理は単純です。
まず、1行めに入力された名前をキューに積みます。
ループでキューが空かをチェックし(1回めは必ず通る)、その名前をキーとしてMemから値を取得し、キーと値を削除します。そして、取得した値をキューに積みます。キーがなかったら(削除できなかったら)キューには積みません。

簡単に具体例を示します。社員を6人(A~F)とします。
入力値はA,B、B,C、C,A、D,E、E,D、F,Cとします。
Memの要素は次のようになります。
A:[C]
B:[A]
C:[B,F]
D:[E]
E:[D]

ここでAを削除します。キューにはCが新たに入ります。
ループの条件チェックでキューにCがあるので2周目に入り、Cを削除します。今度はキューにB,Fが積まれます。
次のループでBを削除し、キューにはAを使いして[F,A]となります。
次のループでFを削除しますがMemにはFをキーとする要素が無いのでキューは変化しません。キューは[A]となっています。
次のループでAを削除しますがAは削除済みなので新たな要素は追加されず、キューは空になります。
結果、Memは次のようになります。
D:[E]
E:[D]

getCount()

社員の数を数えるメソッドです。
Memの全要素の値のSetに含まれる要素数なのでそれを足し合わせて返します。

雑感

この問題は問題を読み終わった時点で頭の中にプログラムが出来上がっていました。ですが、実際の業務でもこのような方法でデータを取り扱うことはあるのでなかなか良い問題だと思いました。